Drop TE cached quantized weights before offloading the training actor

[xiuhu17]

Summary

Drop TransformerEngine's cached quantized weights before offloading the training actor.
Gated by --clear-quantized-weight-workspaces-on-offload (default on).

Memory Optimization for #615 low precision training: nvfp4, mxfp8

Why

Under a low-precision recipe (MXFP8 / NVFP4 / blockwise FP8) with high-precision parameters, TE
quantizes each weight on the fly and caches the result in module._fp8_workspaces. Nothing ever
clears it — the only .clear() in TE is guarded on a recipe-type change — so the cache lives from
the first forward until the process exits.

In colocate RL the actor is offloaded every rollout step, and torch_memory_saver.pause() copies
every tracked allocation to pinned host memory. The workspaces have no region of their own, so they
land in the default tag and get copied too: one D2H per step plus the matching H2D on resume.

That copy is pure overhead. A workspace is a pure function of the high-precision weight, and TE
rebuilds it through the cache-miss path on the next forward. Clearing the dict before
clear_memory() lets the allocator return those segments, so pause() never sees them.

Asserted off under CUDA graphs: a captured graph replays with the workspace address baked in, so
freeing it would let a later allocation reuse that memory.

Expected saving

Per weight element of every quantized module:

recipe	layout	bytes/element
MXFP8, backward_override=None	rowwise + columnwise data (1 B each) + two e8m0 scales (1/32 each)	2.0625
NVFP4, backward_override=dequantized	rowwise packed data (0.5 B) + e4m3 scale, block 16	0.5625

backward_override=dequantized drops the columnwise copy, making NVFP4 in that mode 3.67x cheaper
per element than MXFP8 with the default override.

Measurements

4x B200, colocate, --disable-weights-backuper (so TMS also backs up the param buffer), 5 rollout
steps, paired psutil sampling around pause(). Reported value is host memory moved per offload,
summed over 4 ranks.

GLM-5.2 744B-A40B, 5-layer, NVFP4

Only layer 3's routed experts run NVFP4, so 64 local experts x 2 GroupedLinear = 128 workspaces/rank.
Theory 64 x (2*2048*6144 + 2048*6144) = 2.416e9 elem x 0.5625 B = 1.27 GiB/rank; logged 1.27 GiB.

	mean	stdev
clearing off	58.93 GB	0.77
clearing on	54.30 GB	1.47
saved	4.63 GB	2.00

4 x 1.27 GiB = 5.08 GB expected, 4.63 GB measured — 91% realized, all paired deltas positive.

DeepSeek-V4-Flash, 4-layer, MXFP8

TP=4, EP=4, so 64 local routed experts per layer: 537 workspaces / 12.62 GiB per rank.
Routed experts alone are 4 layers x 64 x 3 x 2048 x 4096 = 6.44e9 elem x 2.0625 B = 12.38 GiB,
98% of the measurement; the 31.9M elem/layer residual is MLA + shared expert + indexer.

	mean	stdev
clearing off	116.04 GB	2.67
clearing on	63.73 GB	4.53
saved	52.30 GB	5.88

4 x 12.62 GiB = 54.20 GB expected, 52.30 GB measured — 96% realized, 45% less host traffic per
offload. The 10x gap versus GLM is the 3.67x per-element cost of MXFP8 plus a much larger
quantized surface (every layer MoE, versus one NVFP4 layer in the GLM config).

Test plan

Both A/B pairs ran 5 steps to Job succeeded. With clearing off, after wake_up reports the full
537 / 12.62 GiB — the workspace really is backed up and restored. With clearing on it reports 0,
and the next step's before offload is back to 537 / 12.62 GiB, confirming the cache-miss rebuild.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.