feat: block-FP8 LoRA colocate support for Qwen3.5-35B-A3B

[mouad-hpc]

Enables block-wise FP8 (e4m3) base + LoRA RL (GRPO) (Megatron trainer ↔ SGLang rollout) for Qwen3.5-35B-A3B.

SGLang's rollout weights are released/resumed via torch_memory_saver, whose resume() discards weight content by design. For a LoRA run the frozen base is served by SGLang from the FP8 checkpoint and is not re-synced, so the resumed block-FP8 base (+ weight_scale_inv) came back as garbage → gibberish rollouts.

• update_weight_from_tensor.py
  • MILES_SKIP_BASE_SYNC skips the redundant frozen-base re-export for LoRA in colocate (extends the existing distributed-LoRA skip); only adapters sync.
  • Gate the SGLang post_process_weights re-quant on the base sync, keeping the call itself as a scheduler barrier for the async CUDA-IPC LoRA handoff.
  • Loop the LoRA load over dtypes (no-op for single-dtype adapters).
• examples/lora/run-qwen3.5-35B-A3B-megatron-moe-lora-fp8.sh — example launcher. Uses --sglang-enable-weights-cpu-backup so the FP8 base survives the colocate release/resume.

Backwards-compatible: all new behavior is gated behind MILES_SKIP_BASE_SYNC (unset → unchanged); the dtype loop is a no-op for single-dtype adapters.

FP8 frozen-base store (trainer-side)

New --fp8-frozen-base-store: after checkpoint load, the frozen LoRA base linear weights are stored as block-fp8 (e4m3 + fp32 scales, 128x128) and the bf16 copy is dropped; a forward pre-hook rematerializes bf16 per layer and the transient is freed again at offload. The quantized base is excluded from the weights backuper (its bf16 is transient) and bridge GEMM-autocast args are forwarded to the provider. GPU-verified on Qwen3.5-35B-A3B MoE-LoRA colocate (TP2/EP8, 12k): 2896 tensors quantized incl. GDN in/out_proj, resident before-offload 10.5 -> 5.8 GB (~45% cut, ~4.7 GB/rank less offload traffic and pinned host RAM), train_rollout_logprob_abs_diff 3.50 == fp8 baseline. Peak is activation-bound and unchanged in effect (+~4.7 GB fp8 buffers during the step); the win is the offloaded/resident state for colocate rollouts.

122B validation (Qwen3.5-122B-A10B, TP2/EP16, 2x8 H200, 12k ctx)

Quantized 1872 tensors, ~9.25 GB/rank freed, trainer resident before-offload ~11 GB (vs ~20 bf16). abs_diff 5.04-5.06 with and without the frozen-base store (fp8 noise at this scale, feature-independent). Also ports upstream radixark#1593 (bridge-LoRA recompute args) — without it the bridge path silently drops --recompute-granularity, which materially inflates long-context peaks.

Per-layer free (--fp8-frozen-base-per-layer-free)

Frees each module's dequantized bf16 as soon as the step no longer needs it: immediately post-forward under no_grad (checkpointed outer forward / forward-only logprob — the recompute pre-hook re-materializes), and via a grad_fn post-hook after the module's dgrad otherwise. Requires activation recompute to pay off (also in this PR: the upstream radixark#1593 bridge recompute-args port). 122B TP2/EP16 verified, 3 steps, flat resident, abs_diff unchanged:

peak GB (122B, 12k)	baseline	chunked512
no store, no recompute	105.2	94.3
store, sleep-free only	OOM	102.9
store + per-layer free + recompute	43.2	32.5

Implementation note: register_multi_grad_hook cannot be used for the deferred free — its closure holds the inputs' grad_fns from hooks inside the same graph, an uncollectible py<->c++ cycle that retained every microbatch graph (~14 GB/step, found via --record-memory-history snapshots).