Add GLM-4 / Kimi 2.5 / OLMo-2, multi-LoRA, on-policy distillation, and more (49 commits)#11
Merged
Conversation
The dummy-dataset fast path in prepare_datasets called PackingDataset on all ranks concurrently, violating the rank-0-first + barrier contract that PR established for _load_or_compute_bins. Non-zero ranks would race rank 0's bin cache write and fail after 10 retries. Gate the dummy path with the same pattern used for real datasets: rank 0 builds the dataset (computing and caching bins), barrier, then non-zero ranks load from cache.
src/xorl/cli/direct_train.py was an unreferenced secondary training entry point. Keeping it alive means every change to training internals risks drifting between it and the primary xorl.cli.train path.
* Raise NotImplementedError instead of silently falling back to HuggingFace When an arch was not in the xorl registry, the loader used to hand the model to AutoModel/AutoModelForCausalLM with only a warning. If an xorl modeling module failed to import, that produced a silent fallback in which no xorl code paths ran, confusing downstream debugging (see). The registry now records per-module import errors, and the loader raises NotImplementedError (surfacing those errors) rather than falling back. * Apply ruff-format
…inistic accumulation
* [Feat] MoE backend improvements: configurable activation, fused GEMM, deterministic accumulation
Three improvements to all MoE backends (native, triton, quack) that
improve correctness, performance, and generality:
1. **Configurable activation function** — thread `act_fn` from MoEExperts
through all backends and EP paths. Previously hardcoded to F.silu().
Uses compile-friendly `use_gelu_tanh` flag for torch.compile paths.
Backward uses generic autograd grad for non-silu activations.
2. **Fused gate_up GEMM** — single `grouped_mm(x, gate_up_proj)` then
chunk, instead of two separate gate/up GEMMs. ~9% faster, eliminates
bf16 rounding difference between fused and split matmuls.
3. **Post-GEMM routing weights** — apply routing weights (expert_scores)
after the down projection GEMM, not before. Mathematically equivalent
(scalar-per-row commutes with matmul) but produces different bf16
rounding. Critical for models with high-magnitude norm weights.
4. **Deterministic accumulation** (triton/quack local path) — replace
moe_gather's index_add_ with reshape(tokens, top_k, hidden).sum().
5x faster (no GPU atomics), deterministic across runs.
Changes per file:
- experts.py: pass act_fn + gate_up_weight to backend functions
- backend/__init__.py: EP compute wrappers (kwargs→positional for apply())
- backend/native.py: fused GEMM path, post-GEMM routing, reshape-sum
- backend/triton.py, quack.py: pass **kwargs through to inner functions
- ops/moe/triton.py: all 4 autograd Functions updated (forward+backward)
- ops/moe/quack.py: all 5 autograd Functions updated (forward+backward)
Verified:
- Qwen3 MoE (silu): max diff 3e-5 across backends (no regression)
- Gemma4 MoE (gelu_pytorch_tanh): bit-exact across all 3 backends
- EP (2 GPU): triton==quack, both activations, forward+backward OK
* style: apply ruff/ruff-format fixes
Address lint CI failures in test files and merged MoE backend files.
* fix: align XORL_DEBUG_EP quack path with post-GEMM routing; drop gemma4 smoke tests
Quack's _forward_debug applied expert_scores before the down GEMM while the
normal path applies it after. Different bf16 rounding meant debug mode changed
numerics. Move the multiplication to match the normal path.
Also remove tests/test_gemma4_* smoke scripts — they call dist.init_process_group
at import time and break `pytest tests/` collection, and the gemma4 module
they depend on isn't in this repo.
* test: drop remaining EP smoke scripts that break pytest collection
test_qwen3_ep_quick.py and test_ep_debug.py call dist.init_process_group
at import time and fail pytest collection. Remove them along with their
_import_fix helper (only referenced by the deleted files).
* refactor: drop dead moe_act variants
Remove classes/functions that are defined but never selected by the backend
registry or called anywhere: TritonEPGroupGemmMoeAct, TritonMoeExpertsFunctionMoeAct,
triton_moe_forward_moe_act, QuackEPGroupGemmMoeAct, QuackMoeExpertsFunctionMoeAct,
quack_moe_forward_moe_act, and native's _run_experts_moe_act / _gate_up_swiglu.
* refactor(moe): thread act_kind string instead of act_fn callable
Replace fragile `_is_gelu_tanh(act_fn)` name-sniffing with an explicit
`act_kind: str` threaded from MoEExperts through the triton/native/quack
backends and their ops-layer autograd Functions.
- Add `normalize_act_kind` (maps e.g. "gelu_pytorch_tanh" -> "gelu_tanh")
and `check_act_kind_supported` in ops/moe/triton.py as the shared source.
- Each backend declares SUPPORTED_ACT_KINDS and validates at entry —
unsupported activations now raise ValueError instead of silently
falling back to SiLU.
- MoEExperts stores self.act_kind (normalized) alongside self.act_fn
(still used by the eager and LoRA paths).
* refactor(moe): rename act_kind -> hidden_act for config consistency
Matches the upstream HF `config.hidden_act` vocabulary users already know.
The value is still normalized ('gelu_pytorch_tanh' -> 'gelu_tanh') via
`normalize_hidden_act` so backends see a canonical set.
* refactor(moe): drop dead split-path branch in TritonMoeExpertsFunction
MoEExperts always stores weights as a single fused gate_up_proj and passes it
to the local path, so the `else` branch (two separate gate/up GEMMs) was
unreachable. Assert gate_up_weight is set and remove the dead branch in both
forward and backward.
Note: QuackMoeExpertsFunction (and native_ep_compute) still use the
split-path pattern — those are separate refactors since they would need a
new fused GEMM implementation, not just branch removal.
* refactor(moe): make hidden_act explicit in backend shims
Both backend/triton.py and backend/quack.py previously relied on `**kwargs`
as a pass-through channel for hidden_act. The docstrings said kwargs were
"ignored" which was wrong — hidden_act was consumed downstream.
Make hidden_act (and gate_up_weight) explicit parameters. Same change in
the ops-layer `triton_moe_forward` / `quack_moe_forward` entry points.
* refactor(moe): use fused gate_up GEMM in quack and native EP paths
Makes the PR's "fused gate+up GEMM" claim honest for all three backends,
and standardizes the kwarg name on `gate_up_proj` (previously `gate_up_weight`
for the local path, split `gate_proj`/`up_proj` for the quack/native EP paths).
- QuackMoeExpertsFunction: single GEMM on gate_up_proj, fused dgrad/wgrad
in backward (mirrors TritonMoeExpertsFunction).
- QuackEPGroupGemm: takes gate_up_proj directly, fused dgrad/wgrad.
- native_ep_compute: takes gate_up_proj and routes through
_run_experts_grouped_mm's fused branch.
- _quack_ep_fused / _native_ep_fused shims: stop slicing +.contiguous()
copying the fused weight.
- Rename gate_up_weight -> gate_up_proj across callers and autograd
Functions for consistency with MoEExperts.gate_up_proj.
QuackTPMoeExpertsFunction (TP path) still uses split GEMMs and would need
a separate refactor.
* test(ep): update test_adapter_source to match fused gate_up_proj signature
The source-string regression test was asserting on the old pre-fuse
signature (_QuackEPGroupGemm.apply(permute_tokens, cumsum, gate_proj, up_proj,
down_proj, expert_scores)). After making quack EP use a single fused
gate_up_proj GEMM, the call site now uses
(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size,
expert_scores, hidden_act). Update the expected substrings to match.
Same for the native EP wrapper.
---------
* set default grad chkpting * linter
* Add sglang_shared_outer LoRA export format Teach save_lora_checkpoint to emit hybrid-shared MoE LoRA directly in SGLang's shared_outer layout (stacked 3D tensors, w1/w2/w3 slots, out-first dim order), and extend the PEFT load path to accept the same keys so roundtrip works. Avoids the external _lora_convert/convert.py re-pack step when targeting SGLang. * Apply ruff-format to lora utils * Wire LoRA export format through server checkpoint save * Export MoE LoRA in PEFT orientation by default * Fix hybrid-shared MoE LoRA export for SGLang * Mirror moe_hybrid_shared_lora and force collective gather for MoE LoRA - Write `moe_hybrid_shared_lora: True` into adapter_config.json when `lora_export_format="sglang_shared_outer"`. SGLang's lora_manager classifies adapters by this flag, and shared_outer is hybrid_shared on disk; without it SGLang mis-classifies as per_expert and rejects loading under `--lora-moe-format hybrid_shared`. - Skip the fast adapter-manager LoRA save path for any MoE LoRA. The fast path reads rank-local params without an EP all_gather and would export only num_local_experts (e.g. 32) instead of num_experts (e.g. 128), silently truncating the saved adapter. --------- Co-authored-by: Ashwinee Panda <apanda@together.ai>
* Forward lr_min and lr_decay_ratio to linear LR schedule build_lr_scheduler dropped lr_min and lr_decay_ratio when lr_decay_style was "linear", so the linear schedule always decayed to ~0 over the full training horizon, regardless of the configured floor or decay ratio. Pass both through, and update get_linear_schedule_with_warmup to honor lr_decay_ratio (matching the cosine path) and to interpolate down to min_lr instead of 0. * Validate LR scheduler args and clean up cosine - Reject lr <= 0 and lr_warmup_ratio outside [0, 1] in build_lr_scheduler; the lambdas divide by init_lr and a non-positive lr would silently NaN. - Drop the dead max(0, factor) clamp and now-redundant assert in the cosine lambda; with min_lr_ratio >= 0 the factor is already non-negative. Hoist lr_decay_steps out of the lambda to match the linear path. - Add tests/optim/test_lr_scheduler.py covering constant/linear/cosine shapes, lr_min floor, lr_decay_ratio behavior, and the new validation.
* Fix stale and dead tests - test_ep_plan_and_shard_tensor: patch get_parallel_state where parallel_plan imports it, not at source - test_ep_group_gemm_propagates_routing_score_gradients[quack]: pass fused gate_up_proj (backend API was unified in); add moe_add_gather to the kernel stub so the test works in isolation - TestBackendGKN.test_backends_match_reference_and_agree: pass fused gate_up_proj (Triton MoE now asserts it is provided) - test_lora_mixed_precision_keeps_base_bf16_and_skips_generic_upcast: patch build_foundation_model / _parallelize on model_builder where they are looked up - Delete tests/ops/test_moe_act.py: _moe_act flag is never read in src; all 36 tests compared identical computations, so they were vacuous - Delete tests/qlora/test_quantize_error_reduction.py: empty placeholder - Delete TestHubAndRemoteDetection in test_shared.py: always skipped because s3fs is not installed * Drop 4-GPU and 8-GPU vocab parallel CE variants test_vocab_parallel_ce_2gpu already exercises the same code path. Removing the larger variants drops the non-e2e suite's max GPU requirement from 8 to 2, letting it fit on a 2-GPU runner with far more concurrency (24 vs 6 slots org-wide). * Add GPU test workflow on self-hosted-h100-2gpu-cu131 Runs pytest -m "not e2e" inside nvcr.io/nvidia/cuda:12.9.1-cudnn-runtime-ubuntu24.04 on the existing org-scoped ARC scale set on turbox. 2 GPUs cover the full non-e2e suite after the vocab_parallel_ce trim. Image choice: CUDA 12.9.1 runtime matches the cu129 torch wheel pinned in pyproject.toml; cudnn-runtime variant includes cuDNN, NCCL, and nvidia-smi; Ubuntu 24.04 ships Python 3.12 which satisfies requires-python. ~3 GB vs ~25 GB for the pytorch container, so cold starts are minutes faster. git is installed at job start because actions/checkout@v4 requires it and the base image is a clean CUDA runtime. * Apply ruff-format to test_moe_gkn_format.py * Install build-essential for Triton JIT The nvcr cudnn-runtime image ships without a C compiler. Triton needs gcc (or cc) to compile the host-side kernel launcher on first use; without it, 108 GPU tests fail with 'Failed to find C compiler'. * Remove pr-test-cpu.yml in favor of GPU workflow The GPU workflow runs pytest -m 'not e2e' which already covers every CPU-marked test. Running them twice (once on dp-cp arc-runner-set and once on turbox H100 pod) wastes H100 time without adding signal. Also drop the arc-runner-set label from actionlint.yaml since no workflow references it anymore. * Restore CPU test workflow and address GPU workflow review CPU lane (arc-runner-set) is restored verbatim so cheap, fast feedback is preserved for PRs that don't need GPU signal — the H100 lane was ~4x slower and contended. GPU workflow review changes: - Switch base image to nvcr.io/nvidia/cuda:12.9.1-cudnn-devel-ubuntu24.04 so gcc/g++/make are present and build-essential apt install is dropped. - Use actions/checkout@v4 and astral-sh/setup-uv@v8 instead of pinned SHAs. - Enable setup-uv caching at /root/.cache/uv. - uv sync --frozen --group test (uv.lock is checked in and verified in sync). - Drop the nvidia-smi debug step. Re-add arc-runner-set to actionlint.yaml since the CPU workflow uses it. * Pin setup-uv to v8.0.0 (no moving v8 tag exists) * Drop uv sync --frozen (uv.lock is gitignored) * Skip cpu-marked tests in GPU lane to avoid duplicate runs CPU lane already runs pytest -m "cpu". With the GPU lane on "not e2e" they double-cover cpu-marked tests on contended H100 capacity. Narrow the GPU lane to "not e2e and not cpu" so each test runs in exactly one lane. * Disable setup-uv GH Actions cache; rely on arc-runner-cache PVC The runner pod already mounts /root/.cache from a PVC (arc-runner-cache), so uv's default cache directory persists across runs without any GH Actions cache round-trip. Keeping enable-cache: true triggered a post-job 'uv cache prune --ci --force' that consistently fails on the ARC k8s container with 'Executing the custom container implementation failed', marking otherwise-green runs as cancelled.
Enforces Conventional Commits on PR titles (which become squash-merge commit messages). Comments the detected version bump on the PR, or fails the check with guidance if the title is invalid. Prerequisite for migrating the release workflow to semantic-release, where commit types drive the version bump.
* Add DeepSeek V4 distributed Muon paths Implements two opt-in features from DeepSeek V4 §3.5.1, both default off: 1. Full-gradient Muon NS (muon_distributed_mode=full_gradient) - Replaces shard-local NS on FSDP2/EP DTensor params with NS on the all-gathered full matrix, recovering the exact Muon update direction. - Momentum/Nesterov stay on the local shard (linear in grad, commutes with sharding) so the optimizer-state buffer stays at local-shard size. - LR adjustment uses the global matrix shape (fixes a bug where adjust_lr_fn saw the local-shard shape under shard_local mode). - Knapsack matrix-to-rank assignment and bounded ZeRO width are not in this PR; every rank in the param's mesh runs NS redundantly. 2. BF16 stochastic-rounded a2a + FP32 local sum (moe_grad_reduce_mode =bf16_a2a_fp32_sum) - Custom FSDPModule.set_custom_reduce_scatter on EP-experts modules: stochastic-round FP32→BF16, dist.all_to_all_single, sum locally in FP32. Halves comm volume vs FP32 reduce-scatter while preserving FP32 accumulation precision. - Stochastic rounding (xorl.optim.stochastic_round) operates on the FP32 bit pattern; unbiased in expectation. Tests: - CPU: stochastic round dtype, unbiasedness, bracket bounds, determinism. - Distributed (2 GPU): full-gradient Muon matches a single-rank oracle; shard-local provably differs on multi-rank. - Distributed (4 GPU): BF16 a2a reduce-scatter matches FP32 reduce-scatter within BF16 ulp bounds, with unbiased mean over many trials. - Distributed (2 GPU): end-to-end FSDP2 hook smoke test. * Address self-review: BF16 hook preconditions, plan typing, 3D EP test - Assert mp_policy.reduce_dtype=fp32 and gradient_divide_factor=1.0 on every expert FSDPModule before installing BF16StochasticAllToAllReduceScatter. Both are correctness preconditions: a non-fp32 reduce_dtype would surface a runtime error during the first backward (far from configuration), and a factor != 1.0 means FSDP applies predivide before the hook and skips the postdivide because the reduce is custom — silently under-weighting grads. Document the same preconditions on the hook's docstring. - Type _MuonUpdatePlan.placements as Optional[Tuple[Placement,...]] and device_mesh as Optional[DeviceMesh] via TYPE_CHECKING imports, replacing the loose Optional[object] / Optional[Tuple] placeholders. - Add a 3D-EP-experts (Shard(1) on the dp mesh) variant to the existing full_gradient distributed test, exercising the deferred-reshape branch in _muon_step that the previous tests only covered for 2D Shard(0). Both full_gradient (matches single-rank oracle) and shard_local (provably differs) layouts are checked. * Apply ruff-format to test files post-merge
* LoRA: merge in fp32 and cast the sum once Current ``LoraLinear.merge_weights`` and ``MoEExpertsLoRA.merge_weights`` round the LoRA delta per element to the weight dtype before adding: W.add_(delta.to(W.dtype)) This loses ~1 weight-dtype ULP per element of Δ. On bf16 weights with MoE top-k routing downstream, tiny per-element drift can flip expert selection at the router (argmax top-k is discontinuous), which cascades through the stack. Empirically, on Qwen3-30B-A3B with random lora_B ~N(0, 0.005) the naive variant degrades K3 by ~20x vs the fp32-sum-then-cast variant: std=0.005 naive K3=1.9e-1 fp32-cast-once K3=1.0e-2 (18x better) The change is a one-liner: upcast W to fp32, add the fp32 Δ, cast the sum back once. W.data.copy_((W.to(fp32) + delta).to(W.dtype)) Same memory as before (result lands in W.dtype). Strictly ≥ precision of the naive variant — bit-exact in fp32, more faithful in bf16/fp16. On realistic trained LoRAs (password-memorization adapters, std << 0.005) both variants produce identical greedy outputs; this change widens the "safe" margin and is the right default. New tests (tests/models/test_lora_merge_fp32_cast_once.py): - zero-LoRA merge is bit-exact - merged weight matches the fp32-sum-then-cast reference bit-for-bit - fp32-cast-once error ≤ naive error (vs true fp32 merged value) - same invariants on MoEExpertsLoRA (fused gate_up + down) All existing LoRA tests still pass (54 total). * style: apply ruff fixes to lora merge test --------- Co-authored-by: Ashwinee Panda <apanda@together.ai>
* Fix MoE EP backward perf regression for train_router=False Two related changes that recover lost throughput on Qwen3.5-style MoE training when train_router=False (the default for ep_dispatch='alltoall' and the only supported setting for ep_dispatch='deepep'): 1. Triton/Quack EP group GEMM backward (src/xorl/ops/moe/triton.py:127, src/xorl/ops/moe/quack.py:316): skip the extra full down-GEMM that computes grad_expert_scores when expert_scores does not require a gradient. With train_router=False MoEBlock detaches routing_weights upstream so ctx.needs_input_grad[5] is False, making the GEMM purely wasted work in backward. train_router=True still computes grads normally. 2. Make routing replay's record/pop of routing_weights opt-in via a new model arg record_routing_weights (default True for safety). When attention forward is deterministic across checkpoint recompute, the regathered routing_weights match the recorded ones, so the cache is unnecessary and disabling it avoids per-MoE-layer pinned CPU allocations + D2H/H2D copies on every step. The arg is threaded through ModelArguments / ServerArguments -> Trainer / ModelRunner -> build_training_model -> build_foundation_model -> config.record_routing_weights -> MoEBlock.from_config. * Switch Muon ns_algorithm default to gram_newton_schulz + batch standard NS The standard_newton_schulz path was the dominant cause of a ~2x training-step regression on Qwen3.5-35B-A3B vs the prior baseline. With grad shape [E, H, I] (E=local experts) it ran an O(E)-deep Python list comprehension calling the upstream 2D _zeropower_via_newtonschulz once per matrix, emitting ~14k extra kernel launches per optimizer step (40 MoE layers x 3 weight pieces x 8 experts x 5 NS steps x ~3 matmuls each). Two complementary fixes: 1. Default the Muon Newton-Schulz backend to gram_newton_schulz in arguments / server_arguments / optimizer factory / Muon.__init__. gram_newton_schulz already batches across experts via baddbmm and recovered the full ~2x throughput in a 1-pod test (tok/s 12k -> 27k, matching the prior baseline's 27,651 tok/s). 2. Add _batched_zeropower_via_newtonschulz that batches the upstream standard NS recurrence with bmm/baddbmm for users who pin ns_algorithm='standard_newton_schulz'. Bit-exact vs the per-matrix loop on representative shapes; ~2.76x faster on a single CUDA microbench at [8, 2048, 768]. The existing test_muon_standard_newton_schulz_preserves_batched_leading_dims test was monkeypatching the per-matrix _zeropower_via_newtonschulz; updated it to monkeypatch the batched variant and check the [B, H, I] flattening contract end-to-end. (cherry picked from commit 130620902de06418c2d180f635338a1143d0e05a) * ci: re-trigger workflows after title update --------- Co-authored-by: Qingyang Wu <qingyang@together.ai>
* Add Cautious Weight Decay (CWD) Implement Cautious Weight Decay from Chen et al. (arXiv:2510.12402): mask the decoupled weight-decay term by I(u_t * x_t >= 0) so decay only acts on coordinates whose optimizer update aligns with the parameter. The original objective is preserved (no implicit regularizer), and the modification is a one-line change with no extra hyperparameters. Plumbed via a top-level `cautious_weight_decay` flag on TrainArgs and ServerTrainArgs. Supported for anyprecision_adamw, signsgd, muon, and adamw (which auto-routes to AnyPrecisionAdamW with fp32 state since the fused torch.optim.AdamW kernel has no per-coordinate decay hook). SGD explicitly rejects the flag. Mask sign proxy chosen per optimizer: - AdamW family: exp_avg (denominator is positive, so sign matches u_t) - SignSGD: grad - Muon: post-Newton-Schulz update tensor (the actual u_t) Tests cover the mask helper, each optimizer's masked path, the cautious==standard equivalence when all signs align, the build_optimizer routing for adamw/signsgd/anyprecision_adamw/muon/sgd, and the post-NS-update mask for Muon. * Address review: rename Muon test to assert hand-computed reference
Co-authored-by: Qingyang Wu <qingyang@together.ai>
* chore(lora): remove unused merge_lora_weights_stacked / unmerge_lora_weights_stacked These two helpers in src/xorl/ops/group_gemm/kernel/lora_utils.py have no in-tree callers and are only re-exported through xorl.lora and xorl.ops.group_gemm.kernel. They date back to the original MoE+LoRA+EP support commit (74be866) and were superseded by MoEExpertsLoRA.merge_weights (the path PR actually patched for fp32-cast-once precision). Removes the function definitions and prunes them from both __init__.py re-export lists. get_lora_delta_weight_stacked is also unused but is left in place for now; flag separately if it should go too. * chore(lora): remove unused stacked LoRA helpers Removes four dead helpers from src/xorl/ops/group_gemm/kernel/lora_utils.py: - merge_lora_weights_stacked - unmerge_lora_weights_stacked - init_lora_weights_stacked - get_lora_delta_weight_stacked None have any in-tree callers (verified by grep across.py /.yaml /.yml / .toml /.json /.md plus check for "import *" forms). They've been dead since the original MoE+LoRA+EP commit (74be866) and were superseded by MoEExpertsLoRA.merge_weights / per-module LoRA paths. The remaining helper, compute_lora_scaling, is kept — it is used by src/xorl/models/layers/moe/lora.py and src/xorl/qlora/modules/moe_experts.py. Both __init__.py re-export lists are pruned accordingly. Risk: external consumers (e.g. SGLang LoRA export, internal tooling) that import these symbols from xorl.lora or xorl.ops.group_gemm.kernel will break their import. Worth a sibling-repo grep before merging. * chore(lora): remove unused stacked LoRA helpers, fold compute_lora_scaling into kernel __init__ Removes four dead helpers from src/xorl/ops/group_gemm/kernel/lora_utils.py: - merge_lora_weights_stacked - unmerge_lora_weights_stacked - init_lora_weights_stacked - get_lora_delta_weight_stacked Verified by repo-wide grep across.py /.ipynb /.yaml /.toml /.json / .md /.sh /.txt plus dynamic-ref and `import *` checks: none of the four have any in-tree caller. They've been dead since the original MoE+LoRA+EP commit (74be866); the live MoE-LoRA path is MoEExpertsLoRA.merge_weights, which is the path actually patched. The only remaining helper, compute_lora_scaling, is folded directly into xorl/ops/group_gemm/kernel/__init__.py and lora_utils.py is deleted — keeping a one-function file for a 5-line scaling helper isn't worth it. Updates the one direct submodule importer (xorl.qlora.modules.moe_experts) to import from the package instead of the now-deleted submodule. The other caller (xorl.models.layers.moe.lora) already imported from the package and needs no change. Risk: external consumers (SGLang LoRA export path, internal tooling) that import any of the four removed symbols from xorl.lora or xorl.ops.group_gemm.kernel will break their import. Cross-repo grep recommended before merging. --------- Co-authored-by: Ashwinee Panda <apanda@together.ai>
Master TCPStore was created with the PyTorch default wait_for_workers=True, which blocks the constructor until world_size-1 workers connect. The design in init_nccl_group is to create a non-blocking listening master first, then fire /init_weights_update_group at the inference endpoints (workers) from a background thread, then complete the NCCL rendezvous in the main thread via _new_process_group_helper. With the default, the master constructor blocks before the background thread can even start, so workers never receive the HTTP request and sync_inference_weights deadlocks until the request times out. Pass wait_for_workers=False so the master returns as soon as it is listening; the actual rendezvous still synchronizes via _init_training_process_group. Update the unit-test fake TCPStore to accept extra kwargs.
* Add OLMo-2 model support Adds Olmo2ForCausalLM under xorl.models.transformers.olmo2, mirroring the Llama 3 / Qwen 3 modules. OLMo-2 differs from Llama in two ways: * Post-norm: layer norms are applied after attention and MLP (post_attention_layernorm, post_feedforward_layernorm); there is no input_layernorm. * Full-axis QK norm: q_norm and k_norm normalize across the entire (num_heads * head_dim) axis prior to reshape, not per-head. Wires the model_type "olmo2" through auto.py so HF configs are loaded into our local Olmo2Config (which carries the TP/PP plans), and adds a checkpoint handler that fuses gate_proj/up_proj into gate_up_proj and q/k/v_proj into qkv_proj on load (and the inverse on save). Includes parity tests against transformers.models.olmo2. * Add tensor-parallel plan for OLMo-2 Wires up TP for OLMo-2 via the standard colwise/rowwise pair, with ``LocalAxisRMSNormShard`` handling the model's full-axis q_norm/k_norm under colwise q/k_proj. Keeps the residual stream Replicate across decoder layers so xorl's loss path (vocab_parallel_cross_entropy matmuls hidden_states @ lm_head.weight.t() directly, bypassing lm_head) sees a full [B, S, H] per rank. * embed_tokens: "embedding" — Replicate output (default). * q/k/v_proj, gate/up_proj: ColwiseParallel() — Replicate input, Shard(-1) output, use_local_output=True. Block internals run on plain local tensors with hidden/tp per rank. * o_proj/down_proj: RowwiseParallel() — Shard(-1) input, Replicate output (all-reduce). Post-norms and residuals see Replicate without any extra plumbing. * q_norm/k_norm: LocalAxisRMSNormShard — weight sharded on dim 0 so each rank's slice matches its colwise q/k_proj output. Computes a local-axis RMS, matching HuggingFace's OLMo-2 reference. * lm_head: ColwiseParallel() — Replicate input, vocab-parallel output. Add tests/distributed/test_olmo2_tp_e2e.py: 2-rank gloo+CPU end-to-end forward + backward through every TP boundary, including lm_head and gradient-flow assertion. Add tests/models/test_olmo2_support.py::test_olmo2_tp_plan_uses_local_axis_qk_norm: structural assertions on the plan (LocalAxisRMSNormShard for q_norm/k_norm, post-norms + final norm not in the plan, vanilla colwise/rowwise elsewhere). Depends on the generic TP infrastructure in (LocalAxisRMSNormShard, DTensor-aware RMSNorm, ParallelStyle resolver passthrough). H100 1-node smoke (8×H100, v_tp2_dp4.yaml): - 15/15 steps complete - ~72k tok/s steady (peaks 73k) - 40.97 GB peak VRAM - loss converging, gradients healthy Closes * Add tensor-parallel plan for OLMo-2 OLMo-2 declares full-axis q_norm/k_norm (over num_heads * head_dim, not per-head). Under colwise q_proj/k_proj, the q/k tensors arrive hidden-sharded, and a full-hidden RMSNorm weight can't be applied directly — so a custom ParallelStyle is needed. Every other model in the repo either has no QK norm or uses per-head + reshape-first, both of which compose with stock ColwiseParallel; the OLMo-2 quirk is scoped to the olmo2/ folder rather than presented as generic TP infra. Pieces: * src/xorl/models/transformers/olmo2/tp_styles.py — new LocalAxisRMSNormShard ParallelStyle: shards a 1-D RMSNorm weight along dim 0 with no input/output redistribute. Each rank's weight slice matches its local q/k slice from colwise q/k_proj. * src/xorl/models/transformers/olmo2/modeling_olmo2.py — new Olmo2QKRMSNorm subclass of RMSNorm. Without TP the parent forward runs unchanged. With TP the forward detects the Shard(0) DTensor weight and runs the fused op directly on local tensors, computing a local-axis RMS that matches HuggingFace's Olmo2RMSNorm reference under TP. Used only for q_norm/k_norm; post-norms and the final norm stay on vanilla RMSNorm. * src/xorl/models/transformers/olmo2/parallelize.py — Replicate- throughout TP plan (default colwise/rowwise everywhere, post-norms and final norm not in the plan). lm_head: ColwiseParallel for use with vocab_parallel_cross_entropy. * src/xorl/distributed/torch_parallelize.py — minimal 3-line ParallelStyle passthrough in _resolve_tp_style so plans can declare style instances directly. Backwards-compatible with string-based plans. Tests (CPU + 2-rank gloo, no GPU required): * tests/distributed/test_olmo2_qk_rms_norm.py — Olmo2QKRMSNorm forward with and without TP, validated against a per-rank local-slice RMSNorm reference under LocalAxisRMSNormShard. * tests/distributed/test_olmo2_tp_e2e.py — full OLMo-2 fwd+bwd through every TP boundary including lm_head and gradient flow. * tests/models/test_olmo2_support.py — TP plan structural assertions (LocalAxisRMSNormShard for q_norm/k_norm, post-norms not in plan, stock colwise/rowwise elsewhere). H100 1-node smoke (8×H100, v_tp2_dp4.yaml — TP=2, FSDP=4, mb=2): - 15/15 steps complete - ~72k tok/s steady, 40.97 GB peak VRAM - loss converging, gradients healthy Closes, --------- Co-authored-by: Qingyang Wu <qingyang@together.ai>
flash_attn-cute pulls in upstream ``quack-kernels`` via ``flash_attn.cute.utils`` during normal training-stack imports. Once ``quack-kernels >= 0.3.11`` is installed, both upstream ``quack`` and the vendored ``xorl.ops.quack`` register ``@torch.library.custom_op`` entries in the **same** ``quack::`` namespace — e.g. both define ``quack::gemm_out`` with mismatched schemas. Whichever module gets imported second silently wins the registration, which produces ``ValueError: vector::reserve`` deep inside the dispatcher when xorl's quack MoE backend invokes the Python wrapper for the schema it expected. This rename moves all vendored op names to ``xorl_quack::`` (gemm_out, gemm_gated_out, gemm_dgated_out, gemm_act_out, gemm_dact_out, gemm_add_out, gemm_add_inplace, gemm_symmetric_out, _softmax_fwd, _softmax_backward, _topk_fwd, _topk_bwd, _rmsnorm_fwd, _rmsnorm_bwd, cross_entropy_fwd_out, cross_entropy_bwd_out). Vendored and upstream now coexist cleanly: ``torch.ops.xorl_quack.*`` for the local copy, ``torch.ops.quack.*`` for the upstream. No public Python API changes — all callers go through ``xorl.ops.quack.gemm_interface.gemm`` etc. Co-authored-by: Ashwinee Panda <apanda@together.ai>
* Add DistSignSGD optimizer
DistSignSGD signs gradients before FSDP2 reduction, turning the
reduce-scatter into a distributed majority-vote update. Complements the
existing per-rank SignSGD by aggregating sign votes across DP workers.
Also wires in two small training-loop improvements:
- build_wandb_init_settings() reads WANDB_INIT_TIMEOUT env var to
extend wandb.init timeout on slow-network runs.
- _log_memory_snapshot no longer issues an NCCL all_reduce during
startup; the snapshot is rank-0 observability only, and the extra
collective could wedge setup before step 1 when ranks progressed
unevenly.
* Address PR review feedback
- count_active_microbatches: build per-mb flags into a single tensor and
issue one batched all-reduce (op=SUM) instead of an all-reduce +.item()
per micro-batch.
- get_distsign_grad_scale_factor: divide by the actual number of (mb,
rank) pairs that cast a sign vote rather than active_microbatches *
dp_size, so ranks with zero valid tokens (sign(0) = 0 abstains) don't
bias the per-step update toward zero.
- distsignsgd: harden the FSDP2 internals walk and assert no FSDP-managed
parameter ends up with the local sign hook installed, preventing
silent double-signing if torch FSDP2 internals shift.
- Revert unrelated changes: _log_memory_snapshot rank-0/no-collective
rewrite and WANDB_INIT_TIMEOUT helper. These are not part of
DistSignSGD; ship separately if needed.
* Address second-round review feedback on PR
- DistSignReduceScatter: force ReduceOp.SUM regardless of the op FSDP
passes through. With reduce_dtype=fp32, FSDP2 may pass AVG (divide by
N), and the trainer's voter-total scale factor already divides by
active_voter_total — inheriting AVG would double-divide and silently
shrink updates.
- sync_sp_gradients: drop the now-unnecessary.to_local() unwrap on
DTensor grads. The pre-PR path called dist.all_reduce on p.grad
directly; the only behavior we need to add for DistSignSGD is the
skip_dtensor_grads early-return.
- _register_local_sign_hooks: raise NotImplementedError when a non-FSDP
DTensor parameter (e.g. a pure TP/PP shard) shows up. Silently
skipping it would let its grads bypass the sign step and pollute the
voter total.
- configure_distsignsgd: switch to direct ps.{dp_replicate_enabled,
ep_enabled, cp_enabled, cp_fsdp_mode, sp_grad_sync_group} access so
an attribute rename surfaces as AttributeError instead of silently
defaulting to False/None and skipping the guard. Also reject
ps.ep_enabled — EP-managed grads use a different fsdp group, so a
single active_voter_total cannot normalize them correctly.
- get_effective_grad_clip_value: clarify in the docstring that under
DistSignSGD the value reported as "grad_norm" is really the L2 norm
of accumulated sign votes (vote_l2_norm), not a true gradient
magnitude.
- Tests: cover the AVG → SUM override, EP rejection, and non-FSDP
DTensor rejection paths; update sync_sp_gradients tests to match
the simplified DTensor handling.
---------
Co-authored-by: Qingyang Wu <qingyang@together.ai>
* Fix FA3 varlen metadata after sync padding * test: cover stale sync-padded cu seqlens
* Add GLM-4 MoE (GLM-4.5/4.6/4.7) training support Port GLM-4 MoE architecture from tomni to xorl, supporting GLM-4.5-Air, GLM-4.5, GLM-4.6, and GLM-4.7 model variants (model_type: glm4_moe). Key architecture features: - Sigmoid-based grouped top-k routing with e_score_correction_bias - Shared experts alongside routed experts - Partial rotary embeddings (partial_rotary_factor) - Optional QK normalization - Dense layers for first_k_dense_replace layers Includes EP/TP parallelization, checkpoint handler for HF weight conversion, and dummy-data benchmark configs (EP=8, EP+Ulysses). Fix _naive_apply_rotary_pos_emb to handle partial rotary automatically when cos/sin dim < head_dim, which was needed for UlyssesAsyncStrategy compatibility with partial-rotary models. * Fix EP plan to use fused gate_up_proj for GLM-4 MoE experts * Fix ruff lint: import sorting and formatting in GLM-4 MoE module Made-with: Cursor * Fix GLM-4.7 MoE routing bug and add model support - Fix e_score_correction_bias routing bug: bias was incorrectly included in routing weights (expert contribution), not just top-k selection. This caused K3 divergence ~10^28 vs SGLang (now K3 ≈ 0). - Add MTP layer weight remapping in checkpoint handler (embed_tokens, norm, lm_head stored under model.layers.92 in HF checkpoint). - Register glm4_moe model type in auto config loader. - Add GLM-4 MoE FLOPs counter with shared expert and dense layer support. - Fix token shifting in non-packing path (packing.py): HF-format labels were not shifted when enable_packing=False. - Minor K3 test script updates for GLM-4.7 compatibility. * Lint cleanup: hoist imports, fix formatting for GLM-4 MoE PR - Hoist `import warnings` to top-level in checkpoint_handler.py - Hoist `get_replay_stage` import to top-level in modeling_glm4_moe.py - Fix ruff format (trailing whitespace alignment) in count_flops.py - Add experiments/ to.gitignore * Fix MoE test failures: propagate moe_act, fused gate_up_proj, skip server tests - base.py: propagate moe_checkpoint_method='moe_act' to MoEExperts modules by setting _moe_act=True during gradient_checkpointing_enable - test_moe_gkn_format.py: pass fused gate_up_proj to TritonMoeExpertsFunction which now requires it for the group_gemm path - conftest.py: auto-skip server-marked e2e tests unless XORL_RUN_SERVER_TESTS=1 to prevent failures in standard pytest runs * Revert packing.py changes from PR The HF-format token shifting fix for SequentialPacker is generic data-pipeline correctness, not GLM-4-specific, and a more complete version of the fix (which also shifts position_ids and handles weights/advantages) is being landed independently on split/trainer-non-packed-shift. * Drop incidental edits from PR:.gitignore, e2e conftest, moe_gkn test, base.py * experiments(k3_tests): raise on out-of-vocab token IDs instead of skipping Silently filtering OOV token IDs masked a real bug — usually SGLang TP padding leaking into the output. Raise ValueError with the same diagnostic message so the failing prompt is the failing prompt rather than a silently-dropped sample. * fix(glm4_moe): use local Glm4MoeRMSNorm with hardcoded fp32 upcast The shared RMSNorm class respects the global rmsnorm_mode (default "native" → F.rms_norm in bf16). On GLM-4 MoE that path interacts pathologically with the triton/quack MoE kernels: K3 KL divergence vs SGLang/HF on GLM-4.5-Air showed a single-token logprob ratio of ~23 (mean K3 0.011, max 0.185) on borderline routing decisions, while every other rmsnorm × moe combination held mean K3 < 0.002. Match HF Glm4MoeRMSNorm: hardcoded fp32 variance + bf16 weight multiply. No global-mode dependency. Other models keep using the shared RMSNorm. * ci: trigger checks * perf(glm4_moe): compile the fp32-upcast RMSNorm forward Glm4MoeRMSNorm intentionally bypasses F.rms_norm to keep an fp32 upcast (the bf16 native path compounds error across GLM-4 MoE's 92+ norm layers). The previous implementation paid the per-call cost of ~5 separate kernel launches × 92 layers ≈ 460 launches per forward pass. Adds a compiled_eager_rms_norm helper to xorl.models.layers. normalization (parallel to compiled_rms_norm and compiled_zero_centered_rms_norm) that wraps the existing eager_rms_norm in torch.compile, fusing the cast/pow/mean/rsqrt/ mul chain into a single Inductor kernel. Glm4MoeRMSNorm.forward now delegates to it. Behaviour unchanged; same fp32 numerics. * style: collapse compiled_eager_rms_norm signature for ruff-format --------- Co-authored-by: Qingyang Wu <qingyang@together.ai>
* Update loss aggregation to normalize across minibatches and ranks * feat: Add reducers support for PPO and IS * refactor: use reducers in model_runner * add SequenceSum * upd surface of losses to return metric tensors instead of item * cleanup * add durations to pytest * metric reduction tests * mean->partial, remove sum * use generalized cu_seq_lens form for seqpartial * remove closure, rename loss -> local_loss_sum
The fixture passed a 2-element `layer_types` pattern alongside `num_hidden_layers=40`. Upstream transformers' `Qwen3MoeConfig` now runs a strict `validate_layer_type` validator (via huggingface_hub dataclasses) that requires `len(layer_types) == num_hidden_layers`, which rejected the short pattern and failed gpu-tests on main. Build the expected per-layer pattern from `full_attention_interval` and assert against the full list. No production code change needed -- real HF configs already provide a full-length `layer_types`.
First batch toward. - distributed/fsdp2/clip_grad_norm.py: narrow `except Exception` around dist.get_world_size to RuntimeError (the only expected failure when the process group hasn't been initialized). - arguments.py: narrow the git-detection except to (subprocess.SubprocessError, FileNotFoundError, OSError); narrow the two get_type_hints sites to (NameError, TypeError, AttributeError) and chain the cause via `from e` so the underlying type-resolution failure isn't lost. - data/prepare/shared.py:515: keep the broad catch (datasets raises a zoo of types) but include type and message in the log so silent "not found" failures are debuggable. Other broad-except sites tracked under.
Two CLI entrypoints (xorl.cli.preprocess and the runner dispatcher in server/runner/setup) print error messages to stdout before falling through to sys.exit(1). Errors belong on stderr — switching to sys.stderr.write keeps the user-visible behavior while letting us drop the per-file T201 ruff ignores PR added (in a follow-up after merges). Refs.
* Start multi-adapter LoRA refactor * Add configurable adapter state restore mode * Add multi-adapter LoRA parity harness * Refactor multi-adapter LoRA server optimizer wiring * Support heterogeneous multi-adapter LoRA sessions * Sync local multi-adapter test and import fixes * Fix server session lifecycle semantics * Add K8s heterogeneous multi-adapter workflow * Add k8s multi-adapter LoRA stress and SGL sync validation * Fix multi-adapter checkpoint and MoE sync edge cases * Fix session checkpoint lifecycle and legacy optim step mapping * Fix adapter eviction and coordinated checkpoint restore * Fix routing replay transport and sampling session tracking * Fix LoRA session save and registration failure paths * Harden adapter restore and checkpoint save failures * Fix weights-only restore LR and kill-session checkpoint paths * Tighten multi-adapter LoRA checkpoint restore validation * Harden multi-adapter LoRA restore coordination * Fix non-packed password memorization evaluation * Harden sampler adapter checkpoint compatibility * Fix merged API type test expectations * Fix pre-commit after main merge * Roll back failed adapter checkpoint loads * Split multi-adapter LoRA foundation * fix: drop stale worker_port reintroduction and update OptimStepRequest test Main removed worker_port from InferenceEndpoint in. This branch predated that and was re-adding it as a required field, breaking add_inference_endpoint and SyncInferenceWeights tests. Drop the worker_port re-addition. Also update test_request_creation_and_serialization to match the new OptimStepRequest API where adam_params defaults to None and learning_rate is the top-level field. * fix: preserve Tinker API compatibility in multi-LoRA foundation * fix: address multi-LoRA foundation review feedback * fix: address PR 265 review regressions * Split multi-adapter LoRA runtime rank substrate * fix: address runtime rank LoRA review feedback * Split multi-adapter LoRA adapter manager core * fix: address adapter manager review feedback * Split multi-adapter LoRA worker coordination * Fix PR 268 worker session coordination regressions * Split multi-adapter LoRA API session lifecycle * Fix API session lifecycle edge cases * Split multi-adapter LoRA checkpoint restore hardening * fix: harden multi-adapter LoRA checkpoint export * Split multi-adapter LoRA sampler inference tracking * fix sampler inference endpoint tracking * fix: route sampler adapter calls to worker port * Split multi-adapter LoRA validation docs and experiments * Fix multi-LoRA stack test expectations * Route register_session through adapter coordinator * Constrain multi-LoRA compare job to nccl nodes * Merge remote-tracking branch 'origin/main' into split/mal-01-foundation * fix: export runtime-rank LoRA state * fix: align adapter coordinator checkpoint loads * Fix API session refresh edge cases * fix: slice multi-adapter lora exports to session rank * fix: detect resolved MoE LoRA targets on save * Fix updated multi-LoRA stack integration * Preserve LoRA dtype for training checkpoints * Document main-merged multi-LoRA validation * Enable deterministic FA3 replay for multi-LoRA tests * Update multi-LoRA validation report * Sanitize multi-LoRA experiment manifests * Serialize IS metrics as scalars * Use namespace-qualified password sampler endpoint * Define password sampler namespace in trainer job * fix: address multi-LoRA review issues * fix: support Kimi tokenizer in multi-LoRA harness * Fix Kimi EP32 multi-LoRA checkpoint restore * Add server full determinism option * fix multi-adapter server review regressions * fix full-weight API compatibility regressions * fix ep-sharded adapter rank0 restore * Validate session spec for EP broadcast restores * chore: apply pre-commit formatting * fix: update register session dispatcher test --------- Co-authored-by: Qingyang Wu <qingyang@together.ai>
… path inference_mode is strictly stronger than no_grad: it additionally skips view tracking and version counters. The eval-only forward path's returned tensors are only consumed by logging and cross-rank reductions that never re-enter autograd, so the strictness is safe here. Sleep/wake_up paths intentionally left on no_grad — they manipulate parameters that subsequently re-enter autograd-tracked code. Refs.
* Fix R3 routing replay SP alignment * Align R3 routing replay with ring attention --------- Co-authored-by: Qingyang Wu <qingyang@together.ai>
* Improve P2P weight sync efficiency * Document P2P weight sync env profile * chore(weight-sync): apply lint formatting * chore: retrigger PR checks * test(weight-sync): update async transfer helper calls * fix(weight-sync): cover qwen fp8 p2p sync gaps * docs(weight-sync): clarify fp8 p2p setup
* Add on-policy distillation MVP
* Add SGLang teacher OPD e2e test
* Add OPD full-rollout e2e test with student/teacher SGLang servers
* Fix nccl_broadcast TCPStore master deadlock
Master TCPStore was created with the PyTorch default wait_for_workers=True,
which blocks the constructor until world_size-1 workers connect. The design
in init_nccl_group is to create a non-blocking listening master first, then
fire /init_weights_update_group at the inference endpoints (workers) from a
background thread, then complete the NCCL rendezvous in the main thread via
_new_process_group_helper. With the default, the master constructor blocks
before the background thread can even start, so workers never receive the
HTTP request and sync_inference_weights deadlocks until the request times
out.
Pass wait_for_workers=False so the master returns as soon as it is listening;
the actual rendezvous still synchronizes via _init_training_process_group.
Update the unit-test fake TCPStore to accept extra kwargs.
* Document NCCL sync hang in OPD full-pipeline test docstring
The test now keeps the disk-based weight refresh path (xorl
save_weights_for_sampler + SGLang update_weights_from_disk) as the
primary validation. Adding a NOTE explaining that the NCCL path was
also wired up in this branch:
- The wait_for_workers=False fix (separate commit) unblocks the
sync_inference_weights call past TCPStore master construction.
- With the fix in place, the broadcast itself now succeeds — the
trainer reaches "Transfer complete: 0.02s, 5 buckets, 26 params".
- The cleanup phase still hangs: SGLang's /destroy_weights_update_group
never returns for the tiny-model setup, which causes
dist.destroy_process_group on the trainer side to hang as well.
That's an SGLang side issue and is out of scope here.
* Add multi-pod OPD k8s pipeline (student SGLang / teacher SGLang / trainer)
Three-pod cross-node validation of the full OPD loop, so the NCCL weight
sync exercises the actual cluster fabric instead of co-locating everything
on the same node:
experiments/local_benchmark/k8s/opd_full_pipeline/
student-sglang.yaml # SGLang sampler, registered as inference endpoint
teacher-sglang.yaml # SGLang with --enable-return-hidden-states
trainer.yaml # xorl.server.launcher + run_opd_pipeline.py
server_config.yaml # 1-GPU bf16 Qwen3-0.6B server config
scripts/opd/
run_opd_pipeline.py # rollout -> teacher hidden -> fwd_bwd -> optim -> sync
submit_k8s_pipeline.sh
Pods publish their endpoints to /shared/opd-coord/<RUN_ID>/{student,teacher}.json.
The trainer driver reads those, registers the student via /add_inference_endpoint,
and runs OPD_NUM_STEPS iterations of forward_backward(opd_loss) + optim_step +
sync_inference_weights, asserting each broadcast succeeds and rollouts diverge.
podAntiAffinity guarantees the three roles land on distinct nodes; nodeAffinity
excludes hosts known to have broken IB / HBM ECC.
* Driver: verify post-sync rollout instead of trusting sync_inference_weights HTTP status
SGLang's /destroy_weights_update_group does not respond after the NCCL
weight broadcast completes (the trainer log reports Transfer complete in
~0.2 s, but the destroy POST stalls at the TCP read), so the trainer's
sync_inference_weights call eventually returns 500. The broadcast itself
has already updated the inference weights, so verify success by sampling
greedy rollouts before and after the sync and checking they diverged.
Also bump the bucket size to 4 GB so all params land in one bucket and
fewer HTTP/NCCL hops are needed per sync.
* k8s OPD: lower mem fraction, raise backoffLimit for tenant admission races
* k8s OPD: drop done marker, keep SGLang pods alive, fix trainer-retry deadlock
* k8s OPD: trainer retries up to 8x and fast-fails on OOM-during-launch
* k8s OPD: longer trainer health probe and tolerate either JSON spacing
* k8s OPD: trainer fast-fails when worker OOMs but launcher survives
* weight-sync: abort training pg on destroy + disable teacher radix cache
Two fixes pulled together:
* nccl_broadcast.destroy_nccl_group: replace
dist.destroy_process_group(self.process_group) with self.process_group.abort().
destroy_process_group calls pg.shutdown(), which cooperatively waits for
the inference peer to finalize. With the matching SGLang fix
(), the inference side now
aborts its comm immediately, so a cooperative shutdown on the trainer
side hangs until the 600 s engine timeout fires. Aborting on both sides
is the right semantic when the broadcast has already completed — the
comm is done, we just want to tear it down.
* opd k8s teacher SGLang: pass --disable-radix-cache. With radix caching
on, /generate's prefill hidden states are returned only for tokens not
already cached, so step 2 of the OPD loop sees a [seq_len-prefix]-shaped
tensor for an input of seq_len tokens and the driver fails. Disabling
the cache forces a fresh prefill per request.
* Driver: assert post-sync rollout diverges from *initial* (not previous step)
After the SGLang+xorl abort fixes the pipeline runs end-to-end, but with
the tiny Qwen3-0.6B model two OPD steps drive lm_head into a degenerate
state that emits the same constant rollout every iteration. The previous
assertion compared post-sync rollout to *the previous step's* post-sync
rollout, which fails when both are degenerate even though the very first
sync demonstrably changed the weights. Snapshot the initial greedy rollout
once and assert that at least one step's post-sync rollout differs from
it.
* lint: ruff format on touched files + whitelist 'NotIn' in codespell
Brings the OPD branch through the repo's pre-commit hooks: ruff-format
reflows a few long lines and codespell stops flagging NotIn (the k8s
nodeAffinity operator).
* k8s OPD: replace personal $HOME paths with /workspace placeholders
tests/test_example_assets.py forbids /home/<user>/ in tracked example and
experiment assets. Switch the OPD manifests to /workspace/home and
/workspace/{xorl,xorl-sglang}, matching the convention
in the existing qwen3-32b-muon-1node-job.yaml and friends. The
home-apanda PVC name is the cluster resource and stays.
* fix: OPD loss reporting — fp32 across SP ranks to keep all_reduce sane
With ulysses sequence sharding, ranks holding only IGNORE_INDEX tokens
hit the OPD early-return paths and produced a bf16 zero-loss while
the rank actually computing the KL returned fp32 from
total_weighted_kl / denom. NCCL's dist.all_reduce silently
reinterprets the bytes when participating tensors disagree on dtype,
turning the per-step reported loss into garbage on the rank-0 reporter
(observed: -1.99e+35 with finite gradients on the compute rank).
Three converging fixes:
- _zero_loss_with_graph casts inputs to fp32 before * 0.0.
- _compute_opd_micro_batch_loss uses fp32 in both the early-return
branch and the total_loss accumulator initializer.
- forward_backward's loss reporter casts loss_report to fp32 before
the cross-rank reduction as a defensive measure.
Also makes max_adapters_per_model environment-overridable
(XORL_MAX_ADAPTERS_PER_MODEL, default 32) so multi-step OPD runs
don't have to evict and trigger SGLang's flaky /unload_lora_adapter.
Validated end-to-end on the Countdown 24-puzzle 1-pod manifest:
16/16 OPD steps clean, loss 1.29 -> 0.58, exact 3/32 -> 8/32.
* review: PP guard, fail-loud cache indices, drop misleading non_blocking
Address review comments on:
- Add a symmetric `if self.pp_enabled: raise NotImplementedError(...)`
inside _compute_opd_micro_batch_loss matching the existing TP guard.
The dispatcher-level guard in _run_forward_backward / _run_forward
already blocks PP, but this prevents a future direct caller from
silently launching OPD under a parallel topology that hasn't been
thought through.
- Replace `flat_indices.clamp(min=0)` in TeacherActivationCache.get
with an explicit IndexError on negative values. The clamp masked
producer bugs (an off-by-one in teacher_cache_indices construction
was caught only via the loss magnitude going wild). Add tests for
both negative and out-of-range index paths.
- Drop `non_blocking=True` on the.to() calls in TeacherHeadManager.get
and TeacherActivationCache.get — the source CPU tensors come from
safetensors / torch.load without `pin_memory()`, so the flag is a
no-op and misleads future readers about the copy behavior.
- Document why TeacherHeadManager keeps only one teacher head on
device at a time (vocab * hidden bytes per teacher; multi-teacher
workloads can revisit with an LRU later).
* Address OPD review comments
* Fix OPD teacher hidden state handling
* feat(weight-sync): add Mooncake P2P backend
* test(weight-sync): add sync stress harness
* fix(weight-sync): fail closed on P2P transfer errors
* fix: correct OPD metrics and shifted teacher fields
* fix(weight-sync): pin mooncake and harden p2p routing
* fix(weight-sync): return measured transfer time
* test(weight-sync): report sync timing breakdown
* fix(models): expand qwen3.5 layer types
* fix(weight-sync): prefer fast endpoint health probe
* test: relax OLMo2 checkpoint tolerance
* test: relax OLMo2 logits tolerance
* fix: expand Qwen3.5 MoE layer types
* fix(weight-sync): stabilize async p2p medium batches
* fix(weight-sync): fail closed on p2p finalization
* fix(weight-sync): clean up p2p diagnostics
* fix(weight-sync): speed up p2p fp8 formatting
* fix(weight-sync): stream fp8 cpu workspace transfers
* feat: add OPD teacher profiling path
* style: format OPD profiling changes
* chore: mark OPD teacher-store helper executable
* chore: make OPD k8s manifests path-generic
* docs: record OPD profile trial results
* docs: fix OPD trial log formatting
* fix(weight-sync): honor remote sync timeout
* test: run OPD weight sync profiling path
* fix: launch multinode workers through python module
* fix: launch workers without torchrun wrapper
* fix(weight-sync): bound fp8 cpu workspace staging
* test: record OPD weight sync trials
* chore(weight-sync): add sparse delta probe
* fix(weight-sync): handle qwen3.5 p2p linear attention slices
* test: record OPD p2p retry after qwen35 sync fix
* chore(weight-sync): mark delta probe executable
* fix(weight-sync): require receiver handles for p2p coalescing
* test: record OPD p2p receiver registration retry
* test: record OPD p2p locator registration retry
* test: record OPD p2p allocator registration retry
* test: record OPD p2p location registration retry
* test: record OPD p2p segment registration retry
* test: record OPD no-expand p2p retry
* fix(weight-sync): support kimi p2p receiver layouts
* Fix P2P weight sync edge cases
* fix(weight-sync): restart cached p2p prepare on fallback
* fix(weight-sync): address p2p fp8 review issues
* fix(weight-sync): handle qwen3.5 p2p linear attention slices
* Record OPD P2P CPU staging retry
* fix server adapter and metric regressions
* fix(weight-sync): complete qwen3.5 p2p opd sync
* Revert "fix server adapter and metric regressions"
This reverts commit fabd8ae3b3b7ccb9af3ac6a079b7e10fc067ea9d.
* fix(server): update adapter, routing, and p2p handling
* fix(opd): correct distributed teacher hidden cache writes
* fix(opd): stabilize fsdp optimizer smoke
Add the OPD profiling pipeline updates used for the async smoke runs, anchor lm_head.forward for FSDP OPD loss, and synchronize after optimizer.step() before releasing gradient storage.
Validation: PYTHONPATH=src PYTHONDONTWRITEBYTECODE=1 pytest -q tests/server/runner/test_opd_runner.py tests/ops/loss/test_opd_loss.py tests/scripts/test_opd_pipeline_payloads.py
* fix(opd): require post-sync version acknowledgement
* fix(opd): support p2p sync for tied lm head
* fix(opd): scope NCCL sync groups per request
* fix(merge): align optim_step LR resolution, IS metrics, and adapter optimizer plumbing
CPU-test fixes after the main→ merge:
- training_ops: keep HEAD's _optim_step_learning_rate helper (registered
sessions fall back to AdamParams default) but extend the chain with
main's _server_default_learning_rate and raise loudly when the session
is not registered, satisfying both the HEAD fallback test and main's
fail-loud test.
- model_runner: restore main's full _accumulate_is_metrics /
_finalize_is_metrics that store Python scalars (via _metric_to_float)
so the IS-only metric tests pass; HEAD's _accumulate_loss_metrics still
drives the OPD path.
- adapters/manager: thread cautious_weight_decay (and the rest of the
manager-level optimizer_config) through the legacy session spec so
optimizer_config={'cautious_weight_decay': True} reaches build_optimizer.
- session_spec: propagate cautious_weight_decay in
session_optimizer_build_kwargs.
- request_processor: only inject routed_experts / routed_expert_logits
for forward_backward so non-MoE backends (e.g. TeacherCacheCPUBackend)
do not get unexpected kwargs.
* fix(merge): always thread routed_experts/logits through model_pass kwargs
The orchestrator request_processor now passes routed_experts and
routed_expert_logits to both forward and forward_backward backend
methods, matching the multi-LoRA replay contract introduced in.
The OPD test backend (TeacherCacheCPUBackend.forward) is updated to
accept the new kwargs so its forward signature is compatible with the
unconditional pass-through.
---------
Co-authored-by: Qingyang Wu <qingyang@together.ai>
…k first * fix(server): R3 routing replay picks Qwen3.6 nested top-k + model_topk first `_extract_topk` now checks `config.text_config.num_experts_per_tok` before `config.num_experts_per_tok` so Qwen3.6 (which nests the field under text_config) does not silently miss it. In `fill_routing_replay`, prefer `self._model_topk` over inferring from `decoded_routing[0][0][0]`. With mixed-width rows (e.g. Qwen3.6 returns some 6-wide and some 8-wide tokens), the previous code would pick row 0's width and crash tensorization on later rows. * chore(lint): apply ruff-format to routing_replay_handler.py
* Add XORL throughput tuner skill * Update XORL throughput tuner repro guidance * Fix renderer provenance block indentation * Add Qwen3.6 benchmark recipe * Add throughput tuner skill tests * feat: add K3 correctness gating artifacts * fix: harden K3 replay scheduling * docs: record Qwen3.5 K3 replay attempt * docs: record Qwen3.5 K3 gate pass * feat(moe): add XORL_MOE_SYNTHETIC_ROUTING=balanced diagnostic knob Adds an env-driven synthetic routing mode that overrides TopKRouter.forward and the K3 replay path with a deterministic balanced top-k pattern (experts cycling [0..top_k-1], [top_k..2*top_k-1],...). Lets the throughput tuner isolate routing-induced load imbalance from real-model training cost. K3 static-trace replay correctly unsets the env so real-traffic traces are not overwritten. * feat(bench): add Q3.5-397B R73/R75 shortctx sweep configs Colocates the 12 qwen3_5_397b_a17b shortctx-8node configs next to the existing R73 K3 gate + shortctx MFU summary under skills/xorl-throughput-tuner/benchmarks/qwen3_5_397b_a17b/configs/. Covers R73 (no-async compile + activation-offload prefetch=4) and R75 (async compile) winning recipes plus R60 ns3, R61 mbs2 offload, R63 alltoall, R69 mbs4, R74 mbs5 prefetch=2 ablations. Also fixes copy-paste bugs from the original sweep generator: r69 wandb_name r67 -> r69, tag nocompile -> compile r73 wandb_name r70/prefetch2 -> r73/prefetch4, tag prefetch2 -> prefetch4 r74 wandb_name r70 -> r74 + _async suffix, tag noasync -> async Supersedes. * fix(skill): address open review comments on throughput tuner Bug: Add qwen3.6-35b-a3b to MODEL_PRESETS so the bundled benchmark recipe hits the EP-divides-experts and ulysses-vs-KV-heads guardrails instead of silently falling through to None. Nits: Producer (run_local_benchmark.py) STEP_RE now captures the peak_mem field the trainer already emits, so benchmark_summary.json carries peak_mem_gb per step; parse_summary in collect_xorl_metrics no longer reports peak_mem_gb_max=None on real runs. discover() preserves log-only sibling runs when other runs in the same sweep tree have benchmark_summary.json (was all-or-nothing per input path). Default use_wandb to false in the portable qwen3_6_35b_a3b benchmark config; the renderer already passes WANDB_MODE/WANDB_ENTITY through for users on clusters that have credentials. Replace brittle 'value: "2"' substring test with a regex check that asserts NCCL_NET_GDR_LEVEL stays as PHB. Make securityContext.privileged a --privileged/--no-privileged flag (default true) so stricter clusters can render without it. Test coverage: Add tests for parse_summary peak_mem plumbing, discover() per-dir preference, the qwen3.6 preset resolution path, and the --no-privileged renderer mode.
* fix(deepep): default to synchronous combine, env-gate unsafe async The MoE output from tokens_post_combine is consumed by the transformer block before the next DeepEP dispatch, so deferring the wait lets downstream compute read incomplete combine data. Force async_combine off in tokens_post_combine unless XORL_DEEPEP_UNSAFE_ASYNC_COMBINE is explicitly set. Keep the API flag for provenance. * Address review: warn on neutered async_combine, cache env, update docs - Cache XORL_DEEPEP_UNSAFE_ASYNC_COMBINE at module import (_ALLOW_UNSAFE_ASYNC_COMBINE) instead of re-reading os.environ per layer call, matching the _DEEPEP_PROFILE pattern. - Log a one-time warning when async_combine=True is silently demoted, so users who set deepep_async_combine in YAML are not left wondering why no overlap occurs. - Update tokens_post_combine docstring to describe the env gate and the underlying default-stream race instead of the original async-path contract. - Update config-reference, expert_parallelism, and moe/deepep docs to flag that the async path is disabled by default and how to opt in. - Mark the guard test as cpu and monkeypatch the cached constant rather than the env. * Address review: simplify env check to == "1" Match the codebase convention used in handler.py for env-var boolean flags.
* fix(qwen3.5): mrope_interleaved must not pairwise-rotate q/k features The interleaved flag controls how T/H/W frequency sections are mixed when MRoPE builds cos/sin, not how q/k features are paired. Match HF/SGLang's standard half-rotate. Previously we rebuilt cos/sin to an interleaved layout and called a pairwise rotate_half, which diverged from the reference implementation on Qwen3.5/Qwen3.6. * Narrow fix to Qwen3.5/3.6 modeling — preserve DSv3 pairwise path The previous revision made `qwen3_5_apply_rotary_pos_emb` ignore its `interleaved` argument outright (`del interleaved`). That silently broke DeepSeek-V3 (`modeling_deepseek_v3.py`), which uses the same shared util with `interleaved=getattr(self.config, "rope_interleave", True)` and depends on the pairwise rotation. Move the fix to the right layer: - Restore the `interleaved=True` branch in the shared util (DSv3 contract). - Drop the `interleaved=getattr(self.config, "mrope_interleaved", False)` kwarg at the two Qwen3.5/3.6 attention call sites, so Qwen always uses half-rotate regardless of `mrope_interleaved`. - Document the convention split (Qwen half-rotate vs. DSv3 pairwise) at the shared util. - Restore pairwise-path tests and add an AST-based regression test that fails if either Qwen3.5/3.6 attention re-introduces an `interleaved=` kwarg on the rotary call.
…py to explicit errors * refactor: convert load-bearing asserts to explicit errors First batch toward. Converts the highest-stakes asserts in arguments.py (parallelism config validation) and data_loader.py (dataset shape and gradient-accumulation invariants) to explicit ValueError/RuntimeError/TypeError. These are runtime-load-bearing checks that python -O would silently strip if left as asserts. Other ~440 asserts across src/xorl/ are not touched in this PR; they need a per-call audit (runtime-check vs dev-invariant). Tracking the rest under. * style: collapse multi-line raise ValueError to satisfy ruff-format
…tter
* fix(fsdp2): canonicalize FSDP-wrapped ReduceOp in BF16 a2a reduce-scatter
FSDP emits a wrapped _ReduceOp (via dist._make_nccl_premul_sum(1.0)) when a
custom gradient_divide_factor is set. The previous SUM/AVG check compared
raw object identity and rejected the wrapped op, blocking configs that mix
BF16StochasticAllToAllReduceScatter with custom divide factors.
Add _canonical_reduce_op() that unwraps the inner RedOpType (SUM, AVG,
PREMUL_SUM-as-SUM since the installer requires factor=1.0) and use it in
both the validation check and the post-reduction divide branch. Add a CPU
unit test that exercises raw RedOpType, wrapped ReduceOp, and the
premul-sum factory.
Not in the q35-397B winning recipe (R45 with this path was below default
throughput) — ship as a robustness/compat fix.
* review: tighten _canonical_reduce_op + mark CPU helper test
- Strengthen PREMUL_SUM comment: call out the install-time
gradient_divide_factor==1.0 check as load-bearing; relaxing it without
inspecting the premul factor here would silently under-weight grads.
- Drop the brittle str(op_type).rsplit('.', 1)[-1] fallback; unknown ops
now fall through to the raw return and are rejected by the downstream
SUM/AVG check.
- Add @pytest.mark.cpu to test_canonical_reduce_op_accepts_fsdp_wrapped_ops
so it's selected by the CPU CI lane (pytest -m cpu); previously only
the module-level distributed mark applied, hiding it from the lane it
was meant to guard.
The previously-pinned commit (5678dd9, Feb 23, 2026) explicitly raises "varlen backward is not yet supported on sm90" — making FA4 unusable for packed-sequence training on H100. The newer commit (59f01d6, May 27, 2026) adds Hopper varlen backward support. Also renames the dep from flash-attn-cute to flash-attn-4 since upstream renamed the package between these commits. Verified on cu129 venv: FA4 varlen + mask_mod forward/backward all match torch SDPA reference within bf16 tolerance. Transitively bumps quack-kernels 0.3.10 -> 0.4.1.
…r cross-node NCCL sync * feat(weight-sync): nccl_simple two-phase backend + sidecar routing for cross-node NCCL sync Adds a cross-node NCCL weight-sync path (trainer FSDP -> sglang TP) that avoids two deadlocks hit on 14B / FSDP=4 / 2-node: 1. nccl_simple backend (two-phase). nccl_broadcast interleaves the FSDP unshard all-gather (intra-node) and the weight-comm broadcast (inter-node) per module; with FSDP ranks 1..N racing ahead of rank 0, the two NCCL communicators enqueue kernels in different orders across ranks and deadlock (observed consistently at the 7th bucket). nccl_simple removes the interleaving: Phase A stages all params to CPU during the FSDP loop (FSDP collectives only), Phase B broadcasts in 1 GiB chunks (weight comm only). Select with --server.sync_inference_method=nccl_simple. 2. NCCL backend cache (XORL_NCCL_BACKEND_CACHE, default on). The NCCL comm init (TCPStore rendezvous + ncclCommInitRankConfig) must happen exactly once; building a new backend each sync re-sends /init_weights_update_group, making sglang attempt a second ncclCommInitRankConfig on the same group name and deadlocking the first comm. Cache the backend per (endpoint, group_name, world_size); initialize() runs once and is reused across syncs. 3. Sidecar-receiver routing (XORL_WEIGHT_SYNC_PORT). Receiving the broadcast inside sglang's NUMA-pinned scheduler starves torch's nonblocking eager NCCL init (device_id=) -- NCCL logs "Init START" but never "Init COMPLETE", while the broadcast kernel never launches. The identical receive in fresh processes runs at ~24 GB/s. This env overrides the receiver port for init/update/destroy so they can target a fresh-process sidecar receiver, while generation and pause/resume keep hitting the real serving port. Also: - HTTP fail-fast + per-endpoint thread logging in nccl_broadcast: if the update HTTP call errors immediately (connection refused), raise before dist.broadcast instead of blocking forever on receivers that were never notified (no more silent 600s hangs). - pause_mode default retract -> in_place (api_types / operations / remote): retract blocked sglang's event loop during the sync. Validated 2-node (trainer FSDP=4 -> sglang TP=4, Qwen2.5-Coder-14B): 31.10 GB / 580 params / 51 buckets synced, ~40s cold (incl. NCCL group init) / ~15s steady-state, repeatable across training steps. Touches weight sync -> run the K3 logprob comparison before merge. * style: apply ruff lint + format fixes --------- Co-authored-by: Ashwinee Panda <apanda@together.ai>
…n3-30B-A3B config) * fix(server): keep forward evals out of inference_mode * Add Qwen3-30B-A3B routed-expert multi-LoRA server config * fix(server): use returned metric ops in forward loop
qywu
force-pushed
the
sync/oss-2026-06-16
branch
from
339012c to
e3f8a3e
Compare
qywu
changed the title
qywu
changed the title
qywu
force-pushed
the
sync/oss-2026-06-16
branch
3 times, most recently
from
d86a335 to
9e89822
Compare
…taNet Replace the stale vendored flash-linear-attention fork with upstream FLA for the GatedDeltaNet kernel + norms, fixing the silent correctness bug. - GDN chunk/recurrent kernels -> fla.ops.gated_delta_rule. The vendored fork ran the Triton gated chunk_bwd_dqkwg, which produces WRONG gradients on Triton>=3.4 / Hopper with no guard; upstream routes that backward to a TileLang kernel. - Norms (RMSNorm / FusedRMSNormGated) -> fla.modules (numerically identical, cos=1.0 fwd+bwd). - Deps: flash-linear-attention pinned to public commit fla-org@97bcb883; tilelang==0.1.11 (first PyPI release with both the b_dq layout fix and Hopper fp32-MMA support needed by the Ulysses CP backward); apache-tvm-ffi==0.1.11. - Kept xorl's pure-torch ShortConvolution + ops/cp CP-halo glue on purpose: upstream's conv is correct but ~3x slower in GDN training (per-call host overhead x ~180 conv calls/step), measured via a same-node A/B. Validated: GPU CI 427 passed; CP-equivalence + Ulysses smoke green; end-to-end loss-match on Qwen3.5-35B-A3B @ 65k.
kiddyboots216
approved these changes
Jun 16, 2026
qywu
force-pushed
the
sync/oss-2026-06-16
branch
2 times, most recently
from
7dff24b to
a29857c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the latest models, training features, and weight-sync improvements — 49 commits, each preserving its original authorship.
Highlights
New model support
Training & optimization
Weight sync & performance
nccl_simpletwo-phase backend + cross-node sidecar routing