feat(metax): FSDP2 full feature parity + Qwen3 training match

[lvyufeng]

Enable FSDP2 (fully_shard + DTensor) on MetaX boxing with complete feature coverage: per-layer wrapping, mixed-precision, optimizers, gradient clipping, state_dict, no-sync accumulation, CPU offload, and 2D mesh (DP×TP).

Verification

FSDP2 features (13/13 pass, ws4):

• per-layer wrapping → all params DTensor, matches single-GPU reference
• mixed-precision bf16 → sharded params kept fp32, tracks single-GPU
• Adam optimizer → foreach path active, matches single-GPU
• clip_grad_norm_ → matches single-GPU
• state_dict → entries are DTensor, reload reproduces trajectory
• gradient accumulation → reduce deferred during no-sync, sum matches per-step reference
• CPU offload → matches single-GPU (the key fix)
• 2D mesh (DP×TP) → builds successfully

Qwen3-0.6B training (20 steps, ws4, batch 2, seq 128):

• Loss trajectory: both 3.170916 → 0.033521
• Worst relative diff: 7.64e-08
• Throughput: flagos 3659.8 tok/s, native 3730.4 tok/s (0.98×)
• Step-0 rank-0 loss bitwise identical (same checkpoint+batches), remaining steps differ in last digits (all_reduce order)

Ops integration: 589 passed, 259 skipped, 1 xfailed, 3 xpassed (9m4s)

Root Cause

Blocking cudaMemcpy (used by FSDP2 CPU offload) only orders against the legacy default stream. FSDP2's reduce-scatter runs in a side stream, so the D2H copy saw zeros. PyTorch creates side streams with cudaStreamNonBlocking, making this deterministic wrong-data (not a race).

Changes

• csrc/aten/copy_ops.cc (+52 lines): SyncCurrentStreamBeforeBlockingCopy() at 7 call sites — after each contiguous() that may enqueue work, before the blocking Memcpy. Guards _to_copy CPU-offload, _copy_from cross-device, and _local_scalar_dense.

• scripts/codegen_ops.py (+37 lines): box plain const Tensor& args in foreach ops (37 kernels) to prevent self-recursion SIGSEGV when unboxed inputs re-dispatch on PrivateUse1.

• torch_fl/comm/process_group.py (+30 lines): _register_inner_backend() so DeviceMesh can read pg.group_name (ProcessGroup::setGroupName forwards only to registered backends).

• Test files (+1078 lines):

  • test_fsdp2_features_metax.py: 13 checks covering all FSDP2 features
  • test_qwen3_fsdp2_metax.py: honest flagos-vs-native comparison (no torch_fl import in native mode, identical checkpoint+batches, pad-masked labels)
  • test_fsdp_live_metax.py: FSDP1+FSDP2 smoke tests (9 checks)

Known Upstream Regression (not caused by this work)

FSDP1 at world_size≥4 is flaky on upstream 11457c9 (profiler commit that enabled real CUDA streams in guard.h):

• ws4: ~50% fail, diverging at step 1
• ws8: produces garbage (1e33, -inf) → memory corruption

Bisection confirms:

• My work on 8695f5a → 5/5 stable
• My work on d530681 → 5/5 stable
• My work on 11457c9 → flaky
• Reverting only my copy_ops.cc → still flaky (not mine)
• Pure upstream 11457c9 → SIGSEGV (needs the codegen fix)

FSDP2 is unaffected (13/13 features, Qwen3 match).

---

🤖 Generated with Claude Code