[diffusion] post_training: run weight update under torch.inference_mode()#31263
[diffusion] post_training: run weight update under torch.inference_mode()#31263kangwangamd wants to merge 6 commits into
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
…de() Model params/buffers are created under torch.inference_mode() (they are 'inference tensors'); the in-place copies in the weight-update path (param.data.copy_, DTensor _local_tensor.copy_, and the layerwise-offload CPU-buffer copies) must therefore run inside an inference-mode context. On torch builds with strict inference-tensor enforcement this otherwise raises 'Inplace update to inference tensor outside InferenceMode is not allowed' and makes /update_weights_from_disk return an HTTP error, failing test_update_weights_from_disk.py. Wrapping _load_weights_into_module in torch.inference_mode() (matching layerwise_offload's own restore path) fixes it with no CUDA-path change.
21483ad to
cac5a71
Compare
|
/tag-and-rerun-ci |
|
@amd-bot ci-status |
CI Status for PR #31263Merge verdict: Caution This PR's changed code is not exercised by any PR-CI test. Changed files: Executed CI failure attribution: AMD: 8 failures (0 related) · NPU: 1 failure (0 related) · Others: 0 independent (2 AMD Executed FailuresAll 8 AMD
Other Executed Failures
(Cascade-only jobs, not counted: pr-test-amd-finish, pr-test-npu-finish, pr-test-extra-finish, pr-test-amd-extra-finish, pr-gate.) Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
The errors are in install deps, which is fixed by #31406 , should update branch and wait for new round of CI. |
6 participants
Motivation
multimodal-gen-test-1-gpushard runningtest_update_weights_from_disk.pyfails with a server-side error and the/update_weights_from_diskAPI returns an HTTP error instead of 200:_apply_weightsthen rolls back (Weight update failed for module 'transformer': Inplace update to inference tensor ...).Root cause
The diffusion pipeline builds/loads model weights while running under
torch.inference_mode(), so the model parameters and buffers are inference tensors. PyTorch forbids in-place mutation of an inference tensor outside an inference-mode context.The weight-update path performs in-place copies. Some are shielded by
.data(which detaches from the inference guard) and some are not:weights_updater.pyload_weights_into_modelparam.data.copy_(...).data)weights_updater.pyload_weights_into_modeldtensor_param._local_tensor.copy_(...)layerwise_offload.pycpu_tensor.copy_(...),cpu_buffer[...].copy_(...)On torch builds with strict inference-tensor enforcement, the unsafe sites raise. This is not device-specific (it reproduces identically on CPU, no accelerator involved) and not a code regression — it is torch-version-sensitive.
Note that
layerwise_offload.pyalready wraps its own restore-path copy intorch.inference_mode(False)+torch.no_grad(), so this pattern is already established in this module.Not ROCm-specific — verified against the NVIDIA lane
test_update_weights_from_disk.pyis part of the sharedrun_suite.py --suite 1-gpu, which runs on both the NVIDIA lane (pr-test.yml,1-gpu-h100) and the AMD lane (pr-test-amd.yml). It is not an AMD-only test.On a recent diffusion PR that exercised the full H100 suite (#31029, 2026-07-14), the NVIDIA
multimodal-gen-test-1-gpujob — including the shard that runstest_update_weights_from_disk— was green (11 passed), with no inference-tensor error. Same sglang code, same day.The difference is purely the torch build each environment ships:
test_update_weights_from_disk1-gpu-h100The failing torch (2.9.0a0 alpha) sits between two versions that both pass, so this is not a simple "newer torch is stricter" case — that specific alpha build enforces the inference-tensor rule where the others don't. NVIDIA isn't passing because of a fix to copy over; it just ships a torch that happens not to trip the latent bug. This change makes the code correct on all torch builds (including that NVIDIA one), which is why it belongs in the shared code path rather than in an image/torch pin.
Modifications
Wrap the body of
_load_weights_into_module()(the single choke point that both the offload and non-offload branches flow through) intorch.inference_mode(), so every in-place copy site — params, DTensor_local_tensor, and the offload CPU buffers — runs inside an inference-mode context. No CUDA-path change; the wrap simply re-enters the same context the tensors were created in.Test
.copy_()on a tensor created undertorch.inference_mode()) and goes away when the copy is wrapped intorch.inference_mode(); verified the updated bytes are correct (torch.allclose).test_update_weights_with_offload_enabled[FLUX.2-klein-base-4B]passes on torch 2.9.1 and 2.11.0 and fails on 2.9.0a0, same sglang commit. NVIDIA h100 (torch 2.11.0) is green for this suite on [Diffusion] post_training: Add LoRA IPC weight sync via lora_merge mode #31029.multimodal-gen-test-1-gpu(test_update_weights_from_disk.py) on CI to confirm the update-weights suite goes green with this change.cc @mickqian @HaiShaw @BBuf @niehen6174
CI States
Latest PR Test (Base): ✅ Run #29537287202
Latest PR Test (Extra): ❌ Run #29537286809