| """Diagnose XPU memory leak by testing each component in isolation. |
| |
| Run on a single XPU tile (no SLURM needed): |
| python tests/diagnose_xpu_leak.py |
| |
| Tests each component for 20 iterations, reporting device_free after each. |
| A decreasing device_free indicates a leak in that component. |
| """ |
| import os, sys, gc |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| import torch |
| import torch.nn.functional as F |
| import torch.distributed as dist |
| from torch.nn.parallel import DistributedDataParallel as DDP |
|
|
| try: |
| import intel_extension_for_pytorch as ipex |
| except ImportError: |
| pass |
|
|
| DEVICE = "xpu:0" if hasattr(torch, "xpu") and torch.xpu.device_count() > 0 else "cuda:0" |
| IMG_SIZE = 128 |
| BATCH = 2 |
| NDIMS = 3 |
| TIMESTEPS = 80 |
| STEPS = 20 |
|
|
|
|
| def get_device_free(): |
| if "xpu" in DEVICE: |
| torch.xpu.synchronize() |
| free, total = torch.xpu.mem_get_info(0) |
| return free / 1024**3 |
| else: |
| torch.cuda.synchronize() |
| free, total = torch.cuda.mem_get_info(0) |
| return free / 1024**3 |
|
|
|
|
| def empty_cache(): |
| gc.collect() |
| if "xpu" in DEVICE: |
| torch.xpu.synchronize() |
| torch.xpu.empty_cache() |
| else: |
| torch.cuda.synchronize() |
| torch.cuda.empty_cache() |
|
|
|
|
| def test_network_forward_only(): |
| """Test: network forward pass without backward or DDP.""" |
| from Diffusion.networks import get_net |
| Net = get_net("om_net") |
| net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE) |
| net.eval() |
| empty_cache() |
| print(f"\n{'='*60}") |
| print("TEST 1: Network forward only (no backward, no DDP)") |
| print(f"{'='*60}") |
| base = get_device_free() |
| print(f" Baseline: {base:.2f} GiB free") |
|
|
| for step in range(STEPS): |
| x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE) |
| with torch.no_grad(): |
| out = net(x=x, y=y, t=t, rec_num=2) |
| del x, y, t, out |
| empty_cache() |
| free = get_device_free() |
| print(f" Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})") |
|
|
| del net |
| empty_cache() |
|
|
|
|
| def test_network_backward_only(): |
| """Test: network forward + backward WITHOUT optimizer.step().""" |
| from Diffusion.networks import get_net |
| Net = get_net("om_net") |
| net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE) |
| net.train() |
| empty_cache() |
| print(f"\n{'='*60}") |
| print("TEST 2a: Network forward + backward ONLY (no optimizer.step)") |
| print(f"{'='*60}") |
| base = get_device_free() |
| print(f" Baseline: {base:.2f} GiB free") |
|
|
| for step in range(STEPS): |
| x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE) |
| out = net(x=x, y=y, t=t, rec_num=2) |
| loss = out.mean() |
| net.zero_grad() |
| loss.backward() |
| del x, y, t, out, loss |
| empty_cache() |
| free = get_device_free() |
| print(f" Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})") |
|
|
| del net |
| empty_cache() |
|
|
|
|
| def test_network_forward_backward(): |
| """Test: network forward + backward with optimizer.step().""" |
| from Diffusion.networks import get_net |
| Net = get_net("om_net") |
| net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE) |
| net.train() |
| optimizer = torch.optim.Adam(net.parameters(), lr=1e-5) |
| empty_cache() |
| print(f"\n{'='*60}") |
| print("TEST 2: Network forward + backward (no DDP)") |
| print(f"{'='*60}") |
| base = get_device_free() |
| print(f" Baseline: {base:.2f} GiB free") |
|
|
| for step in range(STEPS): |
| x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE) |
| out = net(x=x, y=y, t=t, rec_num=2) |
| loss = out.mean() |
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
| del x, y, t, out, loss |
| empty_cache() |
| free = get_device_free() |
| print(f" Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})") |
|
|
| del net, optimizer |
| empty_cache() |
|
|
|
|
| def test_deformddpm_forward_backward(): |
| """Test: full DeformDDPM pipeline (diffuse + recover + backward) without DDP.""" |
| from Diffusion.diffuser import DeformDDPM |
| from Diffusion.networks import get_net, STN |
| Net = get_net("om_net") |
| network = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE) |
| ddpm = DeformDDPM( |
| network=network, |
| n_steps=TIMESTEPS, |
| image_chw=[1, IMG_SIZE, IMG_SIZE, IMG_SIZE], |
| device=DEVICE, |
| batch_size=BATCH, |
| v_scale=5e-5, |
| ).to(DEVICE) |
| ddf_stn = STN(img_sz=IMG_SIZE, ndims=NDIMS, padding_mode="border", device=DEVICE).to(DEVICE) |
| import Diffusion.losses as losses |
| loss_reg = losses.Grad(penalty='l2') |
| loss_ang = losses.MRSE() |
|
|
| optimizer = torch.optim.Adam(ddpm.parameters(), lr=1e-5) |
| empty_cache() |
| print(f"\n{'='*60}") |
| print("TEST 3: DeformDDPM full pipeline (no DDP)") |
| print(f"{'='*60}") |
| base = get_device_free() |
| print(f" Baseline: {base:.2f} GiB free") |
|
|
| for step in range(STEPS): |
| x0 = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE) |
| pre_dvf_I, dvf_I = ddpm(img_org=x0, t=t, cond_imgs=None, mask=1, proc_type=[], text=None) |
| loss_ddf = loss_reg(pre_dvf_I, img=x0) |
| trm_pred = ddf_stn(pre_dvf_I, dvf_I) |
| loss_gen_a = loss_ang(pred=trm_pred, inv_lab=dvf_I, ddf_stn=None, mask=1) |
| loss_tot = 2.0 * loss_gen_a + 4.0 * loss_ddf |
| optimizer.zero_grad() |
| loss_tot.backward() |
| optimizer.step() |
| del x0, t, pre_dvf_I, dvf_I, trm_pred, loss_ddf, loss_gen_a, loss_tot |
| empty_cache() |
| free = get_device_free() |
| print(f" Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})") |
|
|
| del ddpm, ddf_stn, optimizer |
| empty_cache() |
|
|
|
|
| def test_ddp_backward(): |
| """Test: DDP-wrapped network forward + backward (requires init_process_group).""" |
| |
| os.environ.setdefault("MASTER_ADDR", "localhost") |
| os.environ.setdefault("MASTER_PORT", "29500") |
| backend = "ccl" if "xpu" in DEVICE else "nccl" |
| if not dist.is_initialized(): |
| dist.init_process_group(backend=backend, rank=0, world_size=1) |
|
|
| from Diffusion.networks import get_net |
| Net = get_net("om_net") |
| net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE) |
| net = DDP(net, device_ids=[0], find_unused_parameters=True) |
| optimizer = torch.optim.Adam(net.parameters(), lr=1e-5) |
| empty_cache() |
| print(f"\n{'='*60}") |
| print("TEST 4: DDP (find_unused_parameters=True) forward + backward") |
| print(f"{'='*60}") |
| base = get_device_free() |
| print(f" Baseline: {base:.2f} GiB free") |
|
|
| for step in range(STEPS): |
| x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE) |
| t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE) |
| out = net(x=x, y=y, t=t, rec_num=2) |
| loss = out.mean() |
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
| del x, y, t, out, loss |
| empty_cache() |
| free = get_device_free() |
| print(f" Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})") |
|
|
| del net, optimizer |
| if dist.is_initialized(): |
| dist.destroy_process_group() |
| empty_cache() |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--test", type=int, default=0, help="Run specific test (1-4), 0=all") |
| args = parser.parse_args() |
|
|
| torch.xpu.set_device(0) if "xpu" in DEVICE else None |
| print(f"Device: {DEVICE}") |
| print(f"Total memory: {torch.xpu.get_device_properties(0).total_memory / 1024**3:.1f} GiB" if "xpu" in DEVICE else "") |
|
|
| if args.test in (0, 1): |
| test_network_forward_only() |
| if args.test in (0, 2): |
| test_network_backward_only() |
| if args.test in (0, 3): |
| test_network_forward_backward() |
| if args.test in (0, 4): |
| test_deformddpm_forward_backward() |
| if args.test in (0, 5): |
| test_ddp_backward() |
|
|
| print("\nDone. If device_free decreases linearly in any test, that component leaks.") |
|
|
