Native tensor parallelism for wxformer_next (closes #415)#418

Open

jsschreck wants to merge 11 commits into

v2/fsdp2-parallelfrom

v2/native-tp-wxformer

Commits on Jun 11, 2026

models: import NextGenWXFormer from cube-sphere branch
jsschreck
committed
wxformer_next: Linear-projection transformer blocks + conv checkpoint remap
jsschreck
committed
parallel: native DTensor tensor parallelism for _tp_plan models (#415)
jsschreck
committed
tests/manual: tp=2 vs tp=1 parity acceptance run for native TP (#415)
jsschreck
committed
trainers: native-TP-aware checkpointing + gloo TP parity test
jsschreck
committed
parallel: reduced native TP — skip spectral-norm blocks instead of raising
jsschreck
committed
parallel: mixed-mesh-safe gradient clipping for native TP + FSDP2
jsschreck
committed
fix: seed model init identically on all ranks in the gen2 entry point
jsschreck
committed
Merge remote-tracking branch 'origin/v2/fsdp2-parallel' into v2/native-tp-wxformer
jsschreck
committed
review: _tp_plan as an init parameter with the class defaults
jsschreck
committed

Commits on Jun 12, 2026

Merge remote-tracking branch 'origin/v2/fsdp2-parallel' into v2/native-tp-wxformer
jsschreck
committed