[faresobeid]

Activation offloading copies model weights to the CPU, not only activations. With torch.compile, each compiled transformer block saves its graph inputs for backward. For EP experts these include the local tensor of the expert DTensor parameter, a plain torch.Tensor that passes the isinstance(nn.Parameter) check. The offload hook then copies the expert weights to pinned CPU memory and back in every micro batch. On GLM-4.5-Air with EP=8 that is ~553 MB of weights per layer per rank, versus a 64 MB layer input at 64k tokens / CP=8. FSDP already frees and re-gathers that storage, so the copy is pure overhead.

OffloadActivations (and maybe_activation_offloading) now take the model and skip any tensor that shares storage with a model parameter (the local tensor for DTensor parameters). The check only runs for tensors that would otherwise be offloaded. RL and SFT pass the model.

GLM-4.5-Air truncated to 12 layers (real weights), SFT, fake data, 1 node 8xH200, cp=8, ep=8, full AC, max_inflight_activations=5, offload="optimizer", mean of steps 3-6:

seq_len	main, offload on	this PR	offload off
64k: step time / host RAM peak	7.7 s / 458 GB	6.55 s / 393 GB	6.2 s / 369 GB
256k: step time / host RAM peak	43.5 s / 473 GB	42.4 s / 413 GB	41.7 s / 369 GB

Offload host RAM drops from ~90-104 GB to ~24-44 GB per node and the offload step-time overhead drops by 70-75%; the savings scale with the number of MoE layers. Losses and grad norms match main and offload-off within run-to-run noise (~1e-3).

Validation: tests/unit/utils/test_act_offloading.py (gpu): a weight view saved by nn.Linear stays on the GPU. tests/unit/{orchestrator,train,utils,transports} + tests/unit/test_configs.py -m "not gpu" and ruff.

#4238 is stacked on this PR.

---

Note

Medium Risk
Changes activation-offload behavior in the training forward/backward path; fix is localized but affects memory and step time for MoE + compile configs.

Overview
Activation offloading no longer treats parameter storage as activations, which was copying large MoE expert weight tensors to pinned CPU every micro-batch under torch.compile + EP (plain tensor views of DTensor params bypassed the old nn.Parameter check).

OffloadActivations and maybe_activation_offloading now take the model, build a set of parameter storage IDs (using to_local() for DTensor), and skip offload when a saved tensor shares that storage. RL and SFT forward paths pass model into the context manager.

Adds a GPU unit test that weight views saved for backward stay on GPU and are not offloaded.

^{Reviewed by Cursor Bugbot for commit c12bb25. Bugbot is set up for automated code reviews on this repo. Configure here.}