[amarbarosys]

Free constant-folded weights before MIGraphX compiles

Problem

torch.compile(model, backend="migraphx") holds about 3x the model's weight size in device memory at its peak and 2x after compile, counting the eager model:

1. pre_partition_pass constant-folds the weights into GPU copies (e.g. transposed Linear weights) inside each partitioned subgraph.
2. lower_subgraph copies them into the MIGraphX program as literals, then compiles while they are still referenced, so the eager model, the folded copies and MIGraphX's upload are resident at once.
3. After compile, the blocks that held the folded copies stay reserved by PyTorch's caching allocator for the life of the process.

Stage trace, 24x Linear(8192, 8192) fp16 (3072 MiB of weights), device memory in use:

stage	master	this PR
model built	3290 MiB	3290 MiB
after pre_partition_pass	6362 MiB	6362 MiB
MGXModule compile start	6430 MiB	3358 MiB
MGXModule compile end	9552 MiB	6480 MiB

Fix

Once the interpreter has run, every get_attr tensor of the subgraph is a literal in the program. lower_subgraph now drops the interpreter and the subgraph's get_attr references, calls torch.cuda.empty_cache(), then builds the MGXModule. The user's model parameters are untouched (the caller still holds them), so this is not the opt-in deallocate path from #215, which moves them to host and doubles host RSS.

Results

torch 2.13.0+rocm7.2, ROCm 7.2, RX 7900 XTX (gfx1100), Linux; device memory on top of the eager model, two runs each:

	master (e551a86)	this PR
peak during compile	+6217 MiB	+3138 to +3162 MiB
held after compile	+6217 MiB	+3119 to +3145 MiB
output vs eager	max abs diff 0	max abs diff 0
tests/dynamo	354 passed, 10 skipped	355 passed, 10 skipped (incl. new test); see the comment on an intermittent test_reduce_ops_dim failure
tests/dynamo/models (wide_resnet50_2, vit_b_16)	4 passed	4 passed
tests/dynamo wall time (interleaved runs)	66 to 69 s	66 to 69 s

tests/dynamo/quantization (torchao 0.18): 16 passed, 2 failed on both master and this branch. The 2 failures are test_quant_LLM (distilgpt2), which hits transformers DynamicCache not being registered as a pytree during tracing, independent of this change.

Probe script (python mem_probe.py; peak/after numbers above)

Samples used VRAM from amdgpu sysfs every 10 ms during compile. Rerun before opening this PR: master +6209 / +6265 MiB peak, this PR +3177 / +3179 MiB, max abs diff 0.

"""Device memory during torch.compile(backend="migraphx") of a stack of fp16 Linear layers.

Samples the GPU's used VRAM (amdgpu sysfs) every 10 ms while the model compiles and
runs once, and reports the peak and the steady level on top of the eager model.

Usage: python mem_probe.py [layers=24] [width=8192]
"""
import glob
import sys
import threading
import time

import torch
import torch_migraphx  # noqa: F401  (registers the "migraphx" backend)

VRAM_FILES = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used")
MIB = 2**20


def vram_used():
    return max(int(open(f).read()) for f in VRAM_FILES)


def main(layers=24, width=8192):
    with torch.device("cuda"):
        model = torch.nn.Sequential(*[
            torch.nn.Linear(width, width, bias=False, dtype=torch.half)
            for _ in range(layers)
        ]).eval()
        x = torch.randn(8, width, dtype=torch.half)
    weights = sum(p.numel() * p.element_size() for p in model.parameters())
    torch.cuda.synchronize()

    base = vram_used()
    peak = base
    done = threading.Event()

    def sample():
        nonlocal peak
        while not done.is_set():
            peak = max(peak, vram_used())
            time.sleep(0.01)

    sampler = threading.Thread(target=sample)
    sampler.start()
    with torch.no_grad():
        out = torch.compile(model, backend="migraphx")(x)
        torch.cuda.synchronize()
    done.set()
    sampler.join()
    after = vram_used()

    with torch.no_grad():
        diff = (out - model(x)).abs().max().item()
    print(f"weights {weights / MIB:.0f} MiB | on top of eager model: "
          f"peak +{(peak - base) / MIB:.0f} MiB, after compile +{(after - base) / MIB:.0f} MiB | "
          f"max abs diff vs eager {diff:.3g}")


if __name__ == "__main__":
    main(*map(int, sys.argv[1:]))

Test: tests/dynamo/test_lower_subgraph_dynamo.py lowers a subgraph with get_attr weights and asserts the folded constants are gone from the GraphModule and the output matches eager. It fails on master and passes on this branch.

Breaking changes: none. The deallocate option is unchanged.

Related: #215, #218.

[amarbarosys]

Dropped the gc.collect(). No reference cycle holds the folded weights: deleting the interpreter and the subgraph's get_attr attributes frees them by refcount. Checked on CPU with the cyclic GC disabled (0/4 folded tensors alive when the program compiles, 4/4 on master) and on gfx1100, 24x Linear(8192) fp16, three alternated rounds: peak +3141/+3162/+3141 MiB without vs +3162/+3141/+3141 MiB with (master +6213 each), +3141 MiB after compile in every run, output identical. I also corrected the description, which blamed reference cycles for the reserved blocks; it is the caching allocator keeping freed blocks, which empty_cache() releases. On master after compile, torch reports allocated +32 MiB but reserved +3104 MiB above the eager model (3074 MiB cached but free); with this PR, reserved +32 MiB.

Tests on the new head: tests/dynamo (without models/quantization) 354 passed, 1 intermittent failure in test_reduce_ops_dim[dim=0-keepdim=True-op_alias=all.dim]. That test runs the FX path (convert_to_mgx) and never calls lower_subgraph; it also fails intermittently on this PR's first commit, never in my master runs, and failures cluster per process, so I suspect MIGraphX's all/any reduce being sensitive to memory layout rather than this change. tests/dynamo/models 4 passed; tests/dynamo/quantization 16 passed, 2 failed (test_quant_LLM, as on master).

[shivadbhavsar]

There is a deallocate flag you can add to the torch.compile options that does this deallocation during conversion time. Have you tried using that?

[amarbarosys]

Thanks, I had missed it. I measured it on the same setup (24x Linear(8192) fp16, gfx1100, 3 alternated rounds): deallocate=True on master peaks at +3152 to +3162 MiB, the same as this PR (+3150 to +3191 MiB), against +6217 to +6231 MiB on master, and the outputs are identical. So the flag already covers this, and the PR only makes it the default. Should I close this, or would you prefer deallocate to default to True?

[shivadbhavsar]

Thanks, I had missed it. I measured it on the same setup (24x Linear(8192) fp16, gfx1100, 3 alternated rounds): deallocate=True on master peaks at +3152 to +3162 MiB, the same as this PR (+3150 to +3191 MiB), against +6217 to +6231 MiB on master, and the outputs are identical. So the flag already covers this, and the PR only makes it the default. Should I close this, or would you prefer deallocate to default to True?

Great! Im glad that flag works identically. I would like to make deallocate default but some users have run into issues with specific models. Ideally we want to make it more robust for when weights are shared between different modules and other edge cases so that it can be on by default

[amarbarosys]

Makes sense, closing this since deallocate=True already covers it. Thanks!