fix: retain colocated LoRA CUDA-IPC exporter storage

[zhoutong-hai]

Summary

Fix the colocated LoRA CUDA-IPC exporter lifetime in
UpdateWeightFromTensor.

The current full-adapter path releases its exporter tensors and calls
torch.cuda.ipc_collect() immediately after the SGLang load acknowledgement.
The consumer can still be releasing imported handles at that point, so later
allocator reuse can race stale CUDA-IPC handles.

This change:

• keeps the exported adapter payload owned by the updater for the process
  lifetime;
• unloads the previous adapter and waits for the source RPC, propagates a
  source-side failure to the gather group, then barriers and synchronizes CUDA
  before reuse;
• allocates normal CUDA storage once and refreshes it in place for
  --offload-train or PP > 1, where restored/assembled adapter tensors are
  otherwise temporary;
• validates tensor names, shapes, dtypes, and devices before an in-place
  refresh;
• synchronizes the producer stream before serializing reused storage;
• retains the exporter before waiting for the load result; and
• removes eager ipc_collect() / empty_cache() reclamation from this path.

Relation to #1923

This is the same lifecycle invariant and failure family reported in #1923:
keep one exporter allocation stable instead of rotating and reclaiming
CUDA-IPC handles mid-run.

The integration is different enough that this is a current-main
implementation rather than a cherry-pick. #1923 targets the older
kimi-k3 chunked transport with a module-global 512 MiB flat transfer buffer
and a per-chunk acknowledgement. Current main sends one full named-tensor
adapter after PP assembly. This PR therefore retains an exact-layout payload
on each updater and adds explicit unload/error-broadcast/group-barrier
ordering around refresh.

Validation

Focused CPU coverage:

• repeated refresh preserves storage identity and updates values;
• producer-stream synchronization;
• reordered names;
• incompatible names/layouts;
• unload acknowledgement, error propagation, barrier, and CUDA ordering;
• retained payload reuse; and
• retention before the consumer load result.

14 passed, 19 deselected; Ruff, Black, py_compile, and
git diff --check pass.

GPU validation used current main at
778227d6d7cf7b581d1eb07910c873516b6baca9 on 4 nodes / 32 H200s:

• PP4 / TP8 / EP8, all-linear LoRA, CPU offload: three complete
  full-adapter publications (two reuse cycles);
• PP1 / TP8 / EP32, attention LoRA, disk offload: three complete
  publications (two reuse cycles).

Both Ray submissions and both Sky jobs succeeded. All trainer-rank
publications returned successfully. Post-run checks found zero GPU processes,
zero pod restarts, zero cgroup OOM events, and no CUDA-IPC error, Xid, OOM,
NCCL failure, or actor restart.

Follow-up all-linear LoRA parity resolution

This PR only fixes publication/storage lifetime. It does not address
trainer/rollout numerical parity.

The subsequent GLM-5.2 investigation found two independent correctness
problems outside this PR:

• long sparse-DSA prefills were non-repeatable under SGLang's default
  sgl-kernel indexer top-k; deterministic FlashInfer top-k with stable
  tie-breaking made the exact scorer batch repeatable; and
• Megatron fused WGrad accumulation wrote the real shared-outer expert-LoRA
  gradient directly into param.main_grad, bypassing the existing Parameter
  hook and allowing replicated LoRA factors to diverge after the second update.
  That Miles defect is addressed by
  PR #2723.

With deterministic top-k and the optimizer-boundary shared-factor collective,
the complete attention/MLA, dense-MLP, routed-expert gate/up, and routed-expert
down LoRA scope passed an eight-update sustained canary with every publication
returning successfully. The original expert-down_proj hypothesis and
incomplete-mask status are therefore superseded.

PR #2707 remains intentionally scoped to the colocated CUDA-IPC exporter
lifetime; PR #2723 carries the later Miles gradient-synchronization fix.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.