Updating tensorwave fork of axolotl#2
Open
nikhil-tensorwave wants to merge 1405 commits into
Open
Conversation
* Fix fsdp2 sharding. Fix validation of ao version for lr groups * remove validation since axolotl requires ao>0.13.0 already * Move fully_shard of entire module for lora_embedding_A/B out of loop * chore: lint --------- Co-authored-by: bekk02 <ID+bekk02@users.noreply.github.com> Co-authored-by: Wing Lian <wing@axolotl.ai>
…roup in posthog (#3455) * include number of params and rounded est of params so we can easily group in posthog * fix typing
…parallelism (#3444) * fix: correct total_num_steps and batch_size calculation with context parallelism * feat: add test for CP batch size --------- Co-authored-by: NanoCode012 <nano@axolotl.ai>
* mxfp4 axo * import lint * test for qat mxfp4 * config for mxfp4 * add qat: * pass base config * MXFakeQuantizeConfig * lint * tune config so it fits in 32GB VRAM --------- Co-authored-by: Wing Lian <wing@axolotl.ai>
* feat: add sonicmoe * feat: add torch compile for routing * feat: add routing smoke test * feat: add qwen3_5_moe, qwen3_vl_moe, qwen3_omni_moe * fix: disable mlp kernel for sonicmoe too * feat: update to sonicmoe release * chore: update import following new sonicmoe changes * feat: update handling for blackwell * feat: add sonicmoe e2e test * fix: installation for updated sonicmoe * fix: git commit * fix: ignore py req and fix metadata * fix: increase min hidden size to match sonicmoe kernel min * fix: attempt properly interleave and handle unpatch mid-test * chore: refactor teardown better * chore: refactor to re-use rearrange * fix: add idempotency guard * fix: address comments on CI memory and interleave * fix: tests grad, param doublewrapped
…kip ci] * extend pytest-sdist timeout to 30 min for slow/flaky tests * Also preload the cdn cache so it doesn't get stampeded * fix yaml syntax * missing fields * can't pipe to dev/null * Fix nightlies and add 2.10.0 to multi-gpu suite
Co-authored-by: SalmanMohammadi <25081738+SalmanMohammadi@users.noreply.github.com>
* upgrade transformers==5.3.0 trl==0.29.0 kernels * use latest deepspeed fixes * use corect image for cleanup * fix test outputs for tokenizer fixes upstream * fix import: * keep trl at 0.28.0 * handle updated API * use latest trl since 0.28.0 doesn't work with latest transformers * use trl experimental for pad to length * monkeypatch trl with ORPOTrainer so liger doesn't croak * upgrade accelerate * more fixes * move patch for orpotrainer * load the imports later * remove use_logits_to_keep * fix loss_type arg as a list * fetch hf cache from s3 * just manually download the missing model for now * lint for pre-commit update * a few more missing models on disk * fix: loss_type internally now list * fix: remove deprecated code and raise deprecate * fix: remove unneeded blocklist * fix: remove reliance on transformers api to find package available * chore: refactor shim for less sideeffect * fix: silent trl experimental warning --------- Co-authored-by: NanoCode012 <nano@axolotl.ai>
…or parallelism (#3462) [skip ci]
…experts (#3461) * fix: add validation for lora target linear with quantize experts * chore: fix lint * chore: comment * fix: missing link on readme
* add: qwen 3.5 * test for qwen , patch * lint * qwen3 fix on main * Apply suggestions from code review Co-authored-by: NanoCode012 <kevinvong@rocketmail.com> * moe config * config moe * configs and chore * Update examples/qwen3.5/122b-a10b-moe-qlora.yaml Co-authored-by: NanoCode012 <kevinvong@rocketmail.com> * Update examples/qwen3.5/35b-a3b-moe-qlora.yaml Co-authored-by: NanoCode012 <kevinvong@rocketmail.com> * chore for qwen + vlm patch * chore lint * qwen lint * 3_5_moe * Update examples/qwen3.5/README.md --------- Co-authored-by: NanoCode012 <kevinvong@rocketmail.com>
* use new tf32 APIs for torch 2.9+ * also upgrade cce for tf32 fixes and lint
* fix: position_ids casted to int64 for qwen35 patch * fix: to use view instead of reshape to ensure noncontiguous error explicitly * chore: lint
* install flash-linear-attention * handle prequant weights for fsdp2 and ensure loss is not zero * fix type for cu_seqlen, uninstall causal_conv1d * chore: lint * uv pip uninstall doesn't need confirmation
* remove cloud configuration we don't base image for * but we do want it for uv
…_sdpa (#3793) [skip ci]
#3732) * perf(lora): route GatedDeltaNet projections through fused LoRA kernels Qwen3.5's linear-attention (GatedDeltaNet) layers call five LoRA-wrapped projections per layer (in_proj_qkv/z/b/a, out_proj) through the plain peft module forward. peft's _cast_input_dtype round-trips the full [B, S, hidden] activation bf16 -> fp32 (fp32 adapters) -> bf16 (autocast) on every forward, every gradient-checkpoint recompute, and again in backward — numerics-free copy traffic that scales with the activation, not the rank. Route them through axolotl's fused LoRA kernel (the path apply_qkv/apply_o already use), which keeps the activation in its compute dtype and casts only the skinny rank-r adapters: - kernels/lora.py: generic apply_lora_linear() for a single peft-wrapped Linear, through the existing LoRA_O autograd Function (quantized frozen bases via the W_quant slot; bias/dropout/DoRA supported). - monkeypatch/lora_kernels.py: attach apply_<proj> bound methods to every LoRA-wrapped GDN projection, gated on lora_qkv_kernel | lora_o_kernel. Name-matched (in_proj_qkv/...), so qwen3_next (in_proj_qkvz/ba) is untouched; find_linear_attn_in_layer is a no-op on non-GDN models. - qwen3_5/modeling.py: _la_proj_fwd routes the five call sites through the attached helper, with an exact peft fallback when the kernel isn't installed. Numerics-neutral: forward output and grad_X are bit-identical to the peft path; adapter-grad cosine ~1.0. Was-vs-is LoRA SFT on Qwen3.5-0.8B and -2B tracks within the run-to-run nondeterminism floor (step-1 loss bit-identical). Existing lora-kernel patching suite green on non-GDN models (SmolLM2, Qwen3-MoE, gemma). * test(lora): cover fused-LoRA routing of GatedDeltaNet projections Regression tests for the GatedDeltaNet (linear-attention) fused-LoRA routing: - find_linear_attn_in_layer selects Qwen3.5 / Qwen3.5-MoE GatedDeltaNet layers (in_proj_qkv/z/b/a + out_proj) and excludes everything else — most importantly qwen3_next, whose projections are fused as in_proj_qkvz/in_proj_ba and must stay on the plain peft path. Also covers non-GDN layers and partial/missing layouts. - apply_lora_linear matches the plain peft LoRA module forward and grad to bf16 float noise (the routing only removes peft's _cast_input_dtype activation round-trip; the GEMM math is unchanged), including the DoRA branch. Added to the existing lora_kernels patching module rather than a new file: a separate test module collected before test_lora_kernel_patching perturbs the import graph enough that transformers reloads modeling_llama, leaving test_attention_patching_integration asserting against a stale LlamaAttention class object. * refactor(lora): address review — concise docs + per-call GDN patch log - Tighten the apply_lora_linear / _la_proj_fwd docstrings and the LINEAR_ATTN_PROJS comment to a single concise line each. - Replace the module-global one-time _LINEAR_ATTN_PATCH_LOGGED guard (which logged only the first layer's projections, once per process) with a per-call summary: count patched GDN layers and log the projection set with the layer count once after patching. Drops persistent module state and restores visibility across repeated apply_lora_kernel_patches calls. * perf(lora): fuse GatedDeltaNet input projections into one autograd node The GatedDeltaNet input projections (in_proj_qkv/z/b/a) all consume the same post-norm hidden states. Route them through a single fused LoRA autograd Function (LoRA_FusedProj) instead of one per projection: the dropout mask is computed and saved once for the whole group, and the input gradient accumulates in a single buffer rather than N separate dY @ W reductions. With lora_dropout>0 this cuts activation memory substantially (one saved mask/layer instead of four, held across the forward); at dropout=0 it is memory-neutral. Supports bias, dropout, DoRA, and base-only (unadapted) projections. out_proj has a different input and stays on the single-projection kernel. Also: - split LINEAR_ATTN_PROJS into explicit LINEAR_ATTN_IN_PROJS / LINEAR_ATTN_OUT_PROJ (the find_linear_attn selection no longer relies on a positional [:-1] slice). - add integration tests asserting apply_lora_kernel_patches wires the fused in-proj + out_proj onto a real Qwen3.5 GatedDeltaNet layer, plus fused numeric parity / subset-adapter / DoRA tests. * perf(lora): address GDN fused-kernel review feedback - gate GDN in_proj on lora_qkv_kernel and out_proj on lora_o_kernel independently, mirroring self-attn (was: either flag patched both) - accumulate patched-projection names across layers for the INFO log - document the shared-dropout-mask divergence from peft at lora_dropout > 0; trim the new docstrings - drop the dead is-not-None in_proj filter; dedupe the in-proj name tuple to LINEAR_ATTN_IN_PROJS (single source of truth) - add tests: independent gating, backward-under-dropout self-consistency (plain LoRA + DoRA, incl. d_mag) --------- Co-authored-by: Wing Lian <wing@axolotl.ai>
Co-authored-by: Your Name <you@example.com>
…atches (#3778) [skip ci] DeepEP main restructured setup.py and removed the disable_nvshmem code path (deepseek-ai/DeepEP#664), so the two README patches no longer apply (patch failed: setup.py:19). v1.2.1 is the last release they apply to. - Ampere section: drop --depth 1 and add git checkout v1.2.1 (a shallow clone cannot check out the tag) - Notes: state the v1.2.1 pin is required; remove the build-from-HEAD guidance Fixes #3766
…ooleans (#3765) [skip ci] * fix(cli/vllm_serve): let explicit CLI False override config-enabled booleans enable_prefix_caching and enable_reasoning were combined with config values using `cli or cfg`, so a config that enabled either option could not be disabled from the CLI (`False or True` -> `True`). Use explicit None checks, matching the existing enforce_eager precedence, so an explicit CLI value always wins and an omitted CLI value falls back to the config. Adds focused tests for both override and fallback. Fixes #3764 * style: add trailing newline to satisfy end-of-file-fixer pre-commit hook * refactor: trim boolean precedence handling and simplify test Address review: fold the enable_reasoning None fallback into `or False`, shorten the explanatory comment, and collapse the two precedence tests into one parametrized test to reduce bloat. * style: fix trailing newline at EOF (pre-commit) * test fix
#3673) [skip ci] * feat(lora): per-module rank/alpha overrides via rank_pattern / alpha_pattern Adds support for PEFT's rank_pattern and alpha_pattern at training time (new config fields lora_rank_pattern / lora_alpha_pattern) and at merge time (memory-efficient merger now honors them per-module). Also fixes a pre-existing dead check in the AdaLoRA gate (target_rank -> target_r / peft_type == "ADALORA"). * docs(lora): register lora.qmd in _quarto.yml so it builds * fix(lora): mismatch-aware pattern warning on lora_model_dir + rank_pattern merge tests Review feedback from #3673: warning under lora_model_dir now compares cfg patterns against the saved adapter_config.json and only fires on genuine mismatch (generic fallback when the adapter config isn't readable locally); adds end-to-end merge tests for rank_pattern alone and rank+alpha combined; documents resume/continue-training semantics in docs/lora.qmd. * chore: regenerate CLI config options for lora_rank_pattern/lora_alpha_pattern
* fix(lora): stop writing input grads into saved activations in fused LoRA backward LoRA_MLP/LoRA_QKV/LoRA_QK backward aliased the input-gradient buffer onto the saved activation X (torch.mm(..., out=X)). When X is a view created inside another custom autograd Function (e.g. liger RMSNorm returns Y.view(*shape)), dynamo refuses to trace the Function (ApplyTemplateBackward view+inplace error), silently fragmenting decoder layers (gradient checkpointing off) or evicting them from compilation entirely (checkpointing on) under suppress_errors. It also corrupted the saved activation for any second backward with retain_graph. Allocate the input gradient fresh instead. The inplace flag is kept as a no-op for API compatibility. Verified: liger->LoRA_QKV topology compiles fullgraph with zero breaks on torch 2.12; grads match eager; whole-layer graphs restored in both checkpointing modes. * test(lora): cover LoRA_QK and LoRA_MLP in saved-activation regression; drop redundant comment Parameterize the non-mutation test across all three fused kernels changed in the previous commit (each fails on the old aliasing independently) and remove a what-comment flagged in review. Compile-traceability test stays QKV-only: all three share the identical backward code shape. * test(lora): skip CUDA-only kernel tests cleanly on CPU-only hosts
… ci] * docs: document CI/lint/test requirements and e2e-local setup for contributors Document the pinned tool versions CI uses (ruff/ruff-format 0.15.8, mypy 1.19.1, bandit 1.9.4) and add guidance to run pre-commit instead of a system-installed ruff/mypy so local lint output matches CI. Also document the CI test matrix and parallel test invocation, GPU e2e-local setup via the public Docker image (including the cu130 tag for Blackwell), and reference the PR template from the contribution steps. * docs: fix e2e heading level and scope local test command to CPU suite - Address CodeRabbit: bump 'Running e2e tests locally' from h4 to h3 to close the heading-hierarchy gap under Getting Started. - Mirror-CI command excludes tests/e2e (GPU e2e runs in a separate CI job), so the documented local run works on a CPU box. * chore: add Makefile wrapping pinned lint/format/test commands One canonical command per task (make lint/format/test/test-e2e) running the exact CI-pinned tooling, so contributors and agents don't fall back to a system ruff/pytest. Referenced from the contributing guide. * docs: make .pre-commit-config.yaml the single source of truth for lint versions Drop the hardcoded ruff/mypy/bandit version tables from CONTRIBUTING/AGENTS and the literal ruff pin from the Makefile; the Makefile now derives ruff's version from .pre-commit-config.yaml. CI's pre-commit job already reads that file, so versions can no longer drift (the monthly autoupdate bot bumps it). * fix: make 'make test' mirror CI's 'not slow' marker selection The test target claimed to match CI but omitted CI's default -m 'not slow' filter, so local runs could execute a different suite than CI. * docs: condense agent lint/test section, make CI references version-agnostic Review feedback from #3748: AGENTS.md lint/test section trimmed for context economy; CONTRIBUTING no longer hardcodes the Python/PyTorch matrix or Docker tags (tests.yml is the source of truth); make test reworded to approximate CI rather than claim parity. * docs: drop Makefile, point contributor/agent docs at direct commands Per maintainer review on #3748: remove the Makefile wrapper and update CONTRIBUTING.md and AGENTS.md to invoke pre-commit/ruff/pytest directly instead of via make targets.
…n attacks (#3760) [skip ci] * ci: pin all GitHub Actions to full commit SHAs to prevent supply-chain attacks * feat: add exact tags for ref --------- Co-authored-by: XananasX7 <XananasX7@users.noreply.github.com> Co-authored-by: NanoCode012 <nano@axolotl.ai>
* feat: add support matrix * fix: simplify * fix: simplify docs * docs: add support matrix link and remove em-dashes in readme
… [skip ci] * feat(ci): harden workflow checkout credentials + add dependabot * fix(ci): disable setup-uv cache on pypi publish path (cache poisoning) * chore(deps): pin ranged deps to exact versions, drop uv prerelease=allow
* fix(expert_parallel): renormalize gate after capacity drop; restore DeepSpeed validator branches; guard dead tied-keys patch * fix: review pass
* feat(kernels): add NVFP4 LoRA sonicmoe support
* feat(kernels): sonicmoe SM100 NVFP4 W4A4 host driver + correctness gates
Compose quack's GemmGatedSm100/GemmDefaultSm100 with sf_vec_size=16
(blockscaled NVFP4 mainloop + gated epilogue, zero quack kernel changes)
behind our own compile driver, since quack's stock drivers never wire the
two together and guard fp4 off the varlen path.
- libs/sonicmoe/fp4_cute.py: grouped varlen_m engine (expert-sorted packed
rows, cu_seqlens, per-expert weights) + dense composition-smoke path
- libs/sonicmoe/sf_layout.py: blocked SF atom packing, dQaccum-padded
varlen SFA, per-expert per_tensor_scale folding, gate/up interleave
- libs/sonicmoe/nvfp4_quant.py: pure-torch NVFP4 quant/dequant oracle
- benchmarks/sonicmoe_nvfp4/: staged B200 smoke gates (upstream baseline,
dense gated composition, grouped varlen) vs fp32 dequantized oracles
- CPU tests for the host-prep layer (SF layout vs independent formula)
Verified against quack f4f54db0 (v0.5.3) source; GPU validation pending
on the pod via the smoke scripts.
* fix(kernels): sonicmoe fp4_cute fixes from B200 validation
All three smoke gates now pass on B200 (torch 2.10+cu130, quack f4f54db,
cutlass-dsl 4.6.0.dev0, apache-tvm-ffi git):
- call-time GemmDefaultSm100 EpilogueArguments must pass add_to_output=None
(Constexpr field; the namedtuple False default fails the FFI signature)
- fake tensors must be all-symbolic (mixed concrete/sym dims crash varlen
ragged-TMA construction with std::bad_variant_access) and the fp4 K dim
needs an even-divisibility hint for the packed-nibble shape check
- smoke_03 oracle dequantizes with the pts-FOLDED e4m3 scales the kernel
actually consumes; folding cost reported separately (~1e-2 mean rel for
non-power-of-2 pts)
Measured: dense and varlen gated/default all at max_rel ~3.9e-3 vs the fp32
dequantized-operand oracle (pure bf16 output rounding); second forward with
different total_m reuses the compiled kernel.
* feat(kernels): wire sonicmoe fp4_cute W4A4 backend into the NVFP4 MoE-LoRA seam
backend="fp4_cute" now runs end to end: fp4_cute_ops unpacks the torchao
NVFP4Tensor, caches one engine per frozen weight, quantizes activations on
the expert-sorted rows, and wraps the grouped GEMM in an autograd.Function
whose backward does per-expert chunked-dequant dX (never through packed fp4).
_lora_backward_per_group splits dx = g @ W + s*(g @ B) @ A so W_eff is never
materialized and an NVFP4 base dequantizes one expert slice at a time.
_sonicmoe_nvfp4_forward auto-selects fp4_cute on SM100/SM110 when dims align
(AXOLOTL_SONICMOE_NVFP4_BACKEND overrides). smoke_04 gates the seam on a real
NVFP4Tensor: forward plus LoRA/hidden grads vs a pure-torch STE oracle.
* fix(kernels): apply NVFP4 per-tensor scale post-GEMM instead of folding into SFB
Folding pts into the e4m3 block scales truncates: the folded scale is
~amax_block/6, which sits in e4m3's subnormal range (3 mantissa bits at
2^-9 granularity) at real LLM weight magnitudes, costing up to ~11% per
block in smoke tests and more on real checkpoints. Not exact even for
power-of-2 pts once subnormal. The fp4_cute path now keeps SFB at the
stored full-range scales and multiplies the output rows by the per-expert
pts in fp32, which also makes the forward weights identical to the exact
dequant the backward uses.
smoke_04 oracle mirrors the post-scale scheme (forward now matches
bitwise) and gradient checks use a relative-Frobenius metric: impl and
oracle contract the LoRA grads in different orders (dW-first vs z-first),
so elementwise bounds turn bf16 associativity noise into false failures
(verified against fp64 on CPU). All four smokes pass on B200.
* feat(kernels): qwen3_moe NVFP4 adapter + real-checkpoint smoke for the fp4_cute path
Thin clone of the glm_moe_dsa adapter minus the DSA patching: the converter
patterns and the direct expert load already use Qwen3-MoE's checkpoint naming,
so only the model_type gate and registration are new. Matches on either expert
backend (scattermoe or sonicmoe).
smoke_05 reruns the smoke-4 end-to-end comparison on a real MoE layer of
nvidia/Qwen3-30B-A3B-NVFP4, fused through the adapter's own load helpers, and
reports the W4A4-vs-dequant error at real weight magnitudes (OQ1) plus the
fraction of block scales that would go e4m3-subnormal if pts were folded.
* feat(kernels): full-model forward smoke + tiny LoRA train gate for sonicmoe NVFP4
Relax the quantized-training guard under use_sonicmoe too (same frozen
NVFP4 base + LoRA pattern as scattermoe). smoke_06 loads the whole
Qwen3-30B-A3B-NVFP4 through KernelsPlugin/Qwen3MoeAdapter and compares
fp4_cute vs dequant logits (NLL, top-1, KL) on real text.
* fix(kernels): bypass PEFT parametrization merge for sonicmoe expert LoRA
PEFT target_parameters wraps the experts module in ParamWrapper whose
forward materializes base + delta via _activate_lora, which cannot run
on an NVFP4 base (aten.add unimplemented) and defeats the fused low-rank
path. Mirror the scattermoe ParamWrapper fastpath: hand the LoRA tuples
to the base module (_sonicmoe_lora) and call its forward directly.
* feat(kernels): phase 2 driver for sonicmoe NVFP4 quality + shape-sweep benchmarks
200-step LoRA quality runs (fp4_cute vs dequant, shared eval split) plus
12-step throughput/memory probes sweeping context (2k/8k/16k) and
tokens/step (mbs 4), with the scattermoe Marlin W4A16 incumbent.
* perf(kernels): vectorize sonicmoe grouped path with torch._grouped_mm
The LoRA delta/backward, base dX, dense grouped GEMM, and varlen SFA
build all looped over 128 experts in Python with a host sync per
iteration (int(offsets[e]) / tolist), ~24k synced iterations per
training step at 48 layers. On sm90+ these become single-launch
torch._grouped_mm calls (ragged unaligned segments and empty groups
verified exact on B200) and one scatter for the SFA; the loops remain
as the CPU/reference fallback.
* perf(kernels): triton NVFP4 dequant + activation quantizer for the sonicmoe path
torchao's table-gather dequantize dominates the fp4_cute backward (~15ms
per layer at Qwen3-30B shapes) and the ~10-kernel pure-torch activation
quantizer dominates the remaining forward. Both become single Triton
kernels with identical rounding (amax/6 clamp 448, RN e4m3, ties-down
E2M1); torch paths remain the CPU/no-triton fallback.
* fix(kernels): IEEE-RN division in the triton NVFP4 quantizer
Triton's / lowers to reciprocal-multiply; values landing exactly on an
E2M1 bucket boundary flipped vs the torch reference (0.36% of codes).
* perf(kernels): fuse activation quant + swizzled varlen SFA into one triton kernel
The e4m3 scales are written straight to their dQaccum-padded swizzled
addresses, dropping the intermediate scale rows, the scatter, and the
pack/permute chain. Byte-identical output; torch path stays as fallback.
* fix(kernels): unpack all three triton kernels in dequant_nvfp4_triton
* perf(kernels): gather-only route/combine, no atomic index_add in the MoE hot path
Every token appears exactly top_k times in the grouped layout, so the
combine is a gather + fixed-order fp32 sum and the route gather's
backward is the same trick; index_add (atomic-bound, top kernel in the
layer profile) disappears from both directions.
* perf(kernels): one-pass in-place row scaling for the NVFP4 per-tensor scale
Replaces float() -> mul -> to(bf16) (three kernels + an fp32
materialization) with a single bit-identical triton pass.
* perf(kernels): fuse LoRA delta add into the fp4 GEMM epilogue, drop per-layer host syncs
- route_and_group: skip the routed-id guard on CUDA (EP is rejected upstream,
non-EP ids come from a width-E topk) and replace bincount+cumsum with
searchsorted on the already-sorted ids; both host-synced every MoE layer
- LoRA grouped path: stash the contiguous B copy for backward reuse; apply
scaling on the [T, r] intermediates instead of the [T, dim] products
- fp4_cute: compile an add_to_output epilogue variant (constexpr, one callable
cached per value); the LoRA delta is written into the out buffer and the base
GEMM reduce-adds alpha*acc into it, deleting the separate elementwise add and
the pts rowscale pass
- per-expert pts: fold pts_e/pts_ref ratios into the e4m3 SFB (max rel err
5.9e-2 on Qwen3-30B, 0 underflowed blocks) with alpha=pts_ref; backward
dequantizes with the same folded scales so forward and backward weights stay
identical; fold saturation or AXOLOTL_SONICMOE_NVFP4_PTS_FOLD=0 falls back to
stored scales + post-GEMM rowscale (pre-fold numerics)
12-step probe at seq 2048: 0.558 -> 0.498 s/step (+12% tok/s), peak max_active
30.8 -> 28.7 GiB. Smokes 4/5 bitwise vs the folded-scale oracle; 30-step loss
tracks a same-day HEAD control within noise; full-model NLL cost of the fold
+0.017 nats (toggle recovers the old path). Pod suites 89 passed / 1 skipped.
* fix(cli): skip MoE expert fusion in merge-lora when a shard holds a partial expert list
The fuse->merge->unfuse pass in merge_lora_sharded_efficient processes one
shard at a time but assumed each layer's full expert list lives inside that
shard. Checkpoints routinely split expert lists across shard boundaries
(nvidia/Qwen3-30B-A3B-NVFP4 splits layers 12, 26, 40), which either crashed
in torch.cat on mismatched per-projection expert counts or silently fused a
partial subset of experts. Also skip groups whose tensors are still
quantized: stacking raw qdata orphans the per-expert scale siblings.
Expert-count completeness is validated against config.json (num_experts and
variants); incomplete or quantized groups fall through to the per-tensor
pass unchanged.
* docs(kernels): document the NVFP4 sonicmoe backend and correct the quack pin
The validated runtime pin is quack-kernels==0.5.0 from PyPI; f4f54db0 is the
upstream source commit the internals were cross-checked against, not the
installed version.
* perf(kernels): optional fp8 backward weight cache over DeepGEMM for NVFP4 MoE dX
Caches the frozen expert weights once as fp8 e4m3 (128x128-block ue8m0
scales, built from the same folded dequant the forward consumes) and runs
backward dX through DeepGEMM's m-grouped contiguous fp8 GEMM instead of
per-backward whole-weight dequant + bf16 torch._grouped_mm. Microbench on
B200: 1.5-1.8x over the bf16 grouped GEMM at Qwen3-30B backward shapes,
dX rel err ~7e-4 (DeepSeek's validated 1x128 training recipe).
DeepGEMM is an optional source-built dependency; without it the bf16
dequant path is used unchanged. torch._scaled_grouped_mm is deliberately
not in the chain: its rowwise fp8 grouped kernel device-aborts on SM100
(torch 2.10+cu130) and kills the CUDA context. AXOLOTL_SONICMOE_NVFP4_BWD
forces deepgemm or bf16; unset auto-selects.
* perf(kernels): cut fp8 backward host overhead below the GEMM win
The index-plan build (a dozen tiny kernels) and the torch pad+quant chain
cost 3x the fp8 GEMM itself. The plan (dest/m_indices/src) is now cached on
the offsets tensor (up and down projections of a layer share it, and it dies
with the tensor), the pad+quant is one 2D-grid triton kernel that reads only
real rows and rounds bit-identically to per_token_cast_to_fp8 (pow2 scales),
and the contiguous alignment drops to 128 (all padded-layout costs scale
with E * alignment; the GEMM itself also measured slightly faster). B200 at
Qwen3-30B shapes, E=128: dX gate_up 1.96x, down 1.30x over bf16
dequant+grouped_mm with the plan cached.
Smokes 4/5 pin AXOLOTL_SONICMOE_NVFP4_BWD=bf16: they validate the
bf16-backward oracle; smoke 7 owns the fp8 deltas (~4e-2 rel_fro on
dX-dependent grads, matching DeepGEMM's own calc_diff<1e-3 gate).
* fix(kernels): make the fp8 DeepGEMM backward explicit opt-in only
The fp8 weight cache is inherently a full fp8 copy of all expert weights:
+27 GiB persistent on Qwen3-30B for +9% tok/s, which fails the memory
gate as a default. AXOLOTL_SONICMOE_NVFP4_BWD=deepgemm enables it for
VRAM-rich setups; anything else, including unset, keeps the bf16 dequant
backward and never imports deep_gemm.
* feat(kernels): extend the qwen3_moe NVFP4 adapter to qwen3_next
nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4 stores its routed experts in the
same modelopt per-expert layout; the layout inspection already handles the
model's extras (shared expert, linear_attn.out_proj / self_attn.o_proj
NVFP4 linears, per-projection input_scale, the never-loaded mtp.* head).
The conversion mapping is now registered under the model's actual
model_type instead of hardcoded qwen3_moe. Adds the 80B validation config.
* fix(qwen3_next): accept cache_position in the packed GatedDeltaNet patch
The patched decoder layer passes cache_position to linear_attn (matching
current transformers), but the patched GatedDeltaNet forward did not accept
it, so any packed qwen3_next training crashed with an unexpected-kwarg
TypeError on the first linear-attention layer.
* perf(kernels): exact per-row per-tensor scales in the fp4 GEMM epilogue
Replaces the lossy SFB pts-ratio fold as the default: the engine keeps the
stored e4m3 block scales and multiplies the full per-expert scales into the
fp32 accumulator via quack's composable epilogue (ColVecLoad activated by a
compile-time mColVecBroadcast operand + a vec_multiply override on
GemmDefaultSm100; zero quack file changes). Single bf16 rounding at the D
store, composes with the add_to_output fusion, and the backward dequant
returns to plain stored-scales-times-pts, so forward==backward holds
without folded copies. Erases the fold's +0.017 nats full-model NLL cost.
Forward is bitwise against the exact-pts fp32 oracle on both the synthetic
smoke and the real Qwen3-30B checkpoint layer. The old fold stays reachable
via AXOLOTL_SONICMOE_NVFP4_PTS_FOLD=1 for A/B numerics debugging.
* perf(kernels): skip random init of direct-loaded fused expert params
With the routed converters skipped for the direct-load fast path, the fused
experts.gate_up_proj/down_proj land as MISSING keys and transformers'
_initialize_missing_keys fills them with CPU normal_() right before
direct_load_nvfp4_experts overwrites every byte. At Qwen3-Next-80B scale that
is ~155 GB of bf16 randn at a measured 0.11 GB/s, ~23 min of waste per launch
(the actual direct expert load is 48.9s). Mark the fused 3D expert params
_is_hf_initialized before init runs; installed only when the direct load is
armed, from both the qwen3_moe and glm_moe_dsa adapters.
* perf(kernels): fuse up-GEMM + gated activation + LoRA delta into one kernel
Epilogue item (a): GemmGatedAuxAddSm100 subclasses quack's GemmGatedSm100
with a preact-space TileLoad aux (the LoRA delta, added to the fp32
accumulator), the exact per-row pts colvec multiply, and the stock gated
activation, in one blockscaled varlen kernel launch. The up weights stay in
CONCAT [gate; up] layout consumed zero-copy through quack's concat_layout
interleaved view (only the block-scale copy is row-permuted for SFB), so no
packed-weight duplication and the LoRA delta / backward stay in concat space.
The INTERLEAVED preact D is stashed for the swiglu backward.
New GroupedUpProjActLoRA autograd node replaces up-GEMM + separate
gated_activation when backend=fp4_cute, act=silu/swiglu, concat weights, no
bias/limit. Kill switch: AXOLOTL_SONICMOE_NVFP4_FUSED_UP=0. Smokes 04/05 pin
the unfused semantics; smoke 08 owns the fused validation (kernel-level
bitwise preact vs fp32 oracle + module-level fused-vs-unfused grads).
* fix(kernels): fused-up aux must be interleaved in memory, not a concat view
Two GPU findings from smoke 8 on B200: (1) a class-body NamedTuple under
'from __future__ import annotations' stores string annotations that quack's
Constexpr converter cannot eval outside quack's module (NameError: cute);
build EpilogueArguments with the functional NamedTuple form instead.
(2) applying concat_to_interleave to the epilogue TileLoad device-crashes
with an illegal instruction (the epi TMA-load path does not survive the
hierarchical N mode), while the same view on mainloop B is fine and preact
comes out BITWISE vs the fp32 oracle. So concat_layout stays ("B",) only
and the LoRA delta is computed directly in interleaved memory by permuting
the small [2I, r*E] LoRA-B factor rows (no [T, 2I] gather).
Smoke 8 passes: kernel preact bitwise (both token distributions), stage
fused-vs-unfused postact 4.7e-3 rel_fro, module grads 2.4-2.8e-2 (bf16
noise); the 3.3e-2 end-to-end forward diff is the down-GEMM NVFP4
requantization amplifying activation-rounding shifts (bucket flips), bounded
by the stage-level check. Smokes 01-05 and the sonicmoe unit tests pass.
* docs(kernels): document the fused up-projection and its kill switch
* feat(cli): nvfp4 expert-merge writer for fused expert LoRA over per-expert bases
merge-lora previously DROPPED a fused expert LoRA (PEFT ParamWrapper over
experts.gate_up_proj/down_proj) when the base stores experts per-expert
unfused as modelopt NVFP4, since no base key matches the adapter keys.
- _find_param_wrapper_lora accepts both 3D orientations: gpt-oss [E, in, out]
and Qwen3-style [E, out, in] (the latter was rejected, so fused expert
adapters never matched their base params)
- _Nvfp4ExpertMergeWriter claims the per-expert weight/weight_scale/
weight_scale_2 triples of LoRA-targeted expert modules out of each shard
(buffering layers split across shard boundaries), dequantizes via torchao,
fuses to the runtime 3D layout, folds the ParamWrapper delta, then unfuses
and requantizes each expert back to NVFP4 with fresh scales. Output keeps
the base checkpoint layout; under --dequant it emits the merged fused bf16
param instead. Per expert the output is bitwise requant(dequant(base)+delta).
- bases the writer cannot handle (no torchao, non-packed weights) keep the
MERGE INCOMPLETE warning; a packed base without an expert count in
config.json raises instead of silently dropping the adapter
* fix(cli): run the expert-merge writer under no_grad
The PEFT layers built for the delta have requires_grad LoRA weights, so every
emitted tensor dragged a per-layer autograd graph (with its CUDA intermediates)
into the output dict: ~3.6 GiB retained per layer, OOM at 48 layers on a 30B
merge. The other merge paths detach; the writer now disables grad entirely.
* fix(cli): share one nvfp4 outer scale across fused expert group members in the merge writer
Per-member amax recomputation let gate/up weight_scale_2 diverge under the
LoRA delta (ratios up to 4.3x on the real 30B adapter), which the loader's
fuse fold then multiplied into e4m3 block scales already at the 448 max,
overflowing to NaN and poisoning the loaded model (loss 0, grad_norm nan).
* fix(kernels): fold nvfp4 fuse scale ratios toward the max outer scale
Folding proj i's ratio relative to proj 0 casts group scales times a ratio
that can exceed 1 into e4m3, and values past 464 become NaN (the cast only
saturates to 448 below that). Using the elementwise max outer scale keeps
every fold ratio at or below 1, so the cast can never overflow.
* fix(cli): requantize merged nvfp4 weights onto the base checkpoint's grid
Recomputing scales shifts every dequant grid point by the amax drift, so all
elements re-round and a small LoRA delta is replaced by uncorrelated noise:
on the real 30B adapter the emitted delta had cosine ~0.0 with the wanted
delta and attention projections retained ~3-5% of it. Reusing the base block
and outer scales keeps unchanged elements bitwise identical and only moves
codes the delta pushes across a boundary, which preserves the delta's
direction; block scales are bumped only where the merged value outgrows the
original block range. This also inherits gate/up outer-scale equality from
the base, which the loader's fuse step relies on.
Also make the ParamWrapper LoRA finder exhaust all nesting levels in the
exact [E, in, out] orientation before trying the transposed one, so a
transposed outer LoRA cannot shadow an exact inner match (regression from
the orientation fix caught by test_param_wrapper_nesting_dim_filter).
* docs(examples): NVFP4 MoE-LoRA configs for Qwen3-30B-A3B and Qwen3-Next-80B
* chore: cleanup benchmarks
* feat: improve doc for release
* chore: improve doc and cleanup comments
* fix: cleanup nvfp4 example configs
* chore: update base readme
* Update README.md
) The gating bmm in the ScatterMoE forward is autocast-eligible, so under bf16/fp16 autocast the autograd Function output takes the autocast dtype while the saved output_expanded (written by the autocast-unaware Triton kernel from fp32 input) stays fp32. Backward then evaluated fp32 @ bf16 in the d_gates matmul and raised "expected scalar type Float but found BFloat16", crashing bf16 mixed-precision full fine-tuning and LoRA on MoE models. Upcast the branch-local grad_out to output_expanded.dtype at the top of the gates branch in both ParallelLinear and ScatterMoELoRA so the whole gates backward runs in a single dtype; it is a no-op when the dtypes already match.
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
20 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.
Description
Motivation and Context
How has this been tested?
Screenshots (if appropriate)
Types of changes
Social Handles (Optional)