feat(qwen3-8-flash-next): add Engram training (#3690)

* feat(models): add Qwen4-Exp Engram training support

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen4-exp): implement fused sparse QSA training path

Replace the production gathered CUDA attention with a native TileLang forward/backward kernel while retaining a bounded PyTorch oracle for CPU and numerical parity. Keep the discrete QSA router frozen and remove the public main-query chunk knob.

Validated by native kernel job 16304320, layer integration job 16306091, CUDA dispatch job 16307032, and full EP8 4K job 16307225. The final EP8/SGLang statistics payload has SHA256 8777b480e428deec040759039768f6749fef0d1c2db5abb7f90112f48ec26960.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen4-exp): add context parallel training

Add model-owned contiguous context parallelism for QSA, GatedDeltaNet, and PLE, including global routing coordinates, differentiable K/V exchange, raw-token hash boundaries, and the exact nine-token PLE convolution halo. Avoid retaining every decoder hidden state when only fused linear cross entropy needs the final state.

Validated by two-rank CP job 16306590 and matched world128/EP128 100-step jobs 16310033 (CP1) and 16310476 (CP8). W&B run IDs are bdelta-qwen4-tulu3-4k-cp1-world128-ep128-tilelang-v1 and bdelta-qwen4-tulu3-4k-cp8-world128-ep128-tilelang-v1.

CP1 and CP8 used the same data and step segment boundaries. Across 100 logged steps the loss-curve cosine was 0.9999992, mean absolute difference was 0.00054, maximum absolute difference was 0.00267, and final losses were 0.60801 and 0.60800. This records matched curve and checkpoint-payload validation, not uninterrupted optimizer-state resume across earlier segmented runs.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen4-exp): store PLE state as DTensor

Represent each already-local PLE owner shard as one globally shaped Shard(0) DTensor before FSDP snapshots ignored parameters. Keep model-owned All-to-All lookup and gradient scaling local while allowing native DCP model and optimizer state to reshard across owner-world sizes.

Add a same-topology bridge for the earlier owner-v1 optimizer namespace, reject legacy cross-topology loads before payload materialization, and materialize missing Transformer Engine FusedAdam lazy-state destinations with the expected FP32 moments and INT16 master remainder layout.

Validated by four-H100 job 16318843: native world2 save to world4 load/resave, legacy same-topology load, explicit legacy world2-to-world4 rejection, and TE lazy-state checkpoint restoration all passed. Artifact digest: c41302434506a0dfd1b859733c38ef61f81f986f1455fc2c3b225389e4acd00b.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen4-exp): add TileLang THD sparse attention

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Revert "feat(qwen4-exp): add TileLang THD sparse attention"

This reverts commit 7d516d19c2b86879c52e0749ede11b717d5fbed8.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(qwen3-8-flash-next): rename Qwen4-Exp to Qwen3.8-Flash-Next

Rename the model family everywhere: Qwen4Exp* classes to Qwen3_8_FlashNext*,
model_type qwen4_exp(_text) to qwen3_8_flash_next(_text), the
components/models and tests directories, the recipe YAML, env vars, docs
page/slug/nav, and registry entries. No functional changes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(qwen3-8-flash-next): strip legacy checkpoint bridge to a minimal PLE path

The PLE table is now always a globally shaped Shard(0) DTensor in
distributed training, so DCP handles model and optimizer state natively.
Remove the pre-DTensor compatibility machinery:

- stateful_wrappers: drop the legacy owner-checkpoint load bridge and the
  plain-Parameter optimizer-state namespacing; keep TE FusedAdam state
  materialization.
- checkpointing: revert the prepare/finalize load hooks (now identical to
  main).
- owner_sharding: keep only ModelOwnedDTensorSpec(process_group,
  gradient_divisor) and its getter.
- training/utils, parallelizer: remove plain-Parameter owner branches and
  defensive validation.
- qwen3_8_flash_next adapter: keep the zero-copy shard-view mechanism,
  drop alias round-trip verification; single-rank tables carry no contract.

Pre-DTensor optimizer checkpoints can no longer be resumed; resave in the
native DTensor format or reset the optimizer.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(qwen3-8-flash-next): replace owner-sharding spec with a plain grad-divisor attribute

Delete nemo_automodel/shared/owner_sharding.py entirely. The contract is
now one attribute the model sets on its PLE parameter:
weight._nemo_model_owned_grad_divisor = <owner world size>. Presence means
the model owns the sharding (FSDP skips the parameter); the value is the
explicit gradient divisor applied before clipping. Common components read
it with getattr and import nothing.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* revert(checkpoint): drop redundant TE FusedAdam state materialization

TE FusedAdam.step() already lazy-initializes empty state through the same
initialize_state() call, and torch's _init_optim_state dummy step inside
get_optimizer_state_dict() triggers it before DCP planning, so a fresh
optimizer resumes correctly without this helper. Verified against TE 2.14
source and an in-container run. stateful_wrappers.py is now identical to
main.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen3-8-flash-next): THD sparse attention with decoupled KV tokens for packed CP

Restore the TileLang direct-THD sparse GQA kernels (reverted in 6868a9742)
and decouple the KV token count from the query token count: the forward
kernel gains an independent kv_tokens dynamic dimension, the backward and
cast kernels take an independent KV token parameter, and the wrappers
validate route IDs against [0, Tkv). Under context parallelism a rank's
contiguous packed query shard can now attend to the gathered global K/V.

A new H100 test splits a packed sequence into two contiguous query shards
(cutting through a document), runs the kernel per shard against the global
K/V, and checks output/dQ against full-run slices and summed shard dK/dV
against full-run KV gradients.

Validated on 1xH100 (job 16612273): 42 passed covering duplicates,
all-invalid rows, int64 overflow, K=2051, strided inputs, per-document
BSHD parity, padded-gap zeros, the CP shard consistency case, and the
dense TileLang regression file.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen3-8-flash-next): single-GPU packed THD QSA training path

Route packed (THD) batches through the model without context parallelism:

- Indexer: a packed row is unpacked into right-padded per-document rows so
  the existing per-row selection runs unchanged (compression groups restart
  at every document, visibility is bounded by document length); local IDs
  are offset by document starts into global flattened routes.
- QSA layer: with cu_seqlens, CUDA BF16 dispatches to the direct THD
  TileLang kernel on [T, H, D] tensors; CPU uses the differentiable oracle.
- GDN: the decoder layer forwards cu_seqlens so the shared FLA varlen path
  resets the recurrence per document. PLE already resets n-gram context
  after EOS and needs no change.
- Packed + CP remains explicitly guarded (next stage).

CPU parity: packed run equals per-document runs bitwise for routes and
within default tolerances for outputs and input gradients. GPU (1xH100,
job 16613941): 58 passed across the THD kernel suite, packed parity
including a full-size-head THD-vs-dense layer comparison, and the QSA and
dense TileLang regression files.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen3-8-flash-next): packed (THD) context parallelism

Extend contiguous CP to packed rows:

- Qwen3_8_FlashNextCPContext carries optional global_cu_seqlens; the CP
  sharder pops cu_seqlens from the batch, shards the [1, T] tensors
  contiguously, and replicates the boundaries in the context.
- The indexer gathers raw index keys and rotary values globally (frozen, no
  gradients) and routes each document intersecting the local shard with the
  existing per-row selection at its own document offset; CP padding rows
  keep all -1 routes.
- The QSA layer takes boundaries from the CP context and feeds local packed
  queries plus the gathered global K/V to the decoupled THD kernel (CUDA) or
  the oracle (CPU).
- The GDN wrapper stashes the global boundaries and builds its
  BlockdiagCpModelState with the CP padding tail as an isolated segment, so
  the shared FLA blockdiag core resets recurrence per document.

Validation: model suite 103 passed on CPU including a 2-rank gloo packed
CP2-vs-CP1 parity test (bitwise global routes with a document cut by the
shard boundary, outputs and input gradients against CP1 slices, summed
parameter gradients, zero CP padding tail) and a GDN boundary-assembly
test. GPU 1xH100 (job 16615067): 60 passed across the THD kernel, packed,
QSA, and dense TileLang suites.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen3-8-flash-next): resolve legacy qwen4_exp checkpoint identifiers

Released checkpoint dumps are immutable and still store model_type
qwen4_exp(_text) and architecture Qwen4ExpForConditionalGeneration from
before the rename. Register legacy config aliases and an architecture
mapping so AutoConfig and NeMoAutoModelForCausalLM resolve those dumps to
the renamed classes. Also synthesize per-token document IDs for the
inherited packed GDN conv path when a cu_seqlens-only packed batch carries
no attention mask.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(qwen3-8-flash-next): materialize legacy sub-configs and accept loader seq_lens

- Legacy config aliases need explicit pass-through __init__: transformers'
  PreTrainedConfig.__init_subclass__ dataclass-wraps subclasses and replaces
  a missing __init__ with a generated one that leaves text/vision sub-config
  dicts unmaterialized.
- The packed-sequence loader emits seq_lens/seq_lens_padded/qkv_format
  rather than cu_seqlens. The model forward and the CP sharder now convert
  those to physical slot boundaries (packed_boundaries_from_seq_lens), so
  packer-introduced intra-document padding stays inside its document
  identically for CP1 and CP-sharded runs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen3-8-flash-next): advertise packed CP capability on the TileLang backend

The recipe gate supports_cp_with_sequence_packing rejected packed CP because
the model declared no packed attention backend. Declare
_packed_cp_attn_backends = ("tilelang",): the fused THD QSA path owns packed
routing, and other attention backends stay rejected.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* docs(qwen3-8-flash-next): use the Hub model ID and hybridep dispatcher in the recipe

Point the recipe and docs at Qwen/Qwen3.8-Flash-Next on the Hugging Face Hub
instead of a local checkpoint env var, drop the oc.env interpolations to
match the other examples, switch the expert dispatcher to hybridep, and
reduce the slurm walltime to 01:00:00.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(qwen3-8-flash-next): topology-independent packing and real-length boundaries

Packed CP1-vs-CP8 parity showed mismatched tokens per step: the THD packer
takes cp_size, which rounds per-document padded lengths and changes pack
composition per topology, and the boundary glue read seq_lens_padded even
though the packer lays tokens out contiguously by real length.

- The recipe loader now packs with cp_size=1 for models that own their CP
  attention: model-owned CP shards the packed row contiguously, so
  per-document CP padding is a TE blockdiag concern that must not leak into
  pack composition.
- packed_boundaries_from_seq_lens builds physical boundaries from the REAL
  seq_lens and appends the trailing pack padding as its own segment, so pad
  tokens never join a real document; seq_lens_padded is ignored.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* docs(qwen3-8-flash-next): record packed CP validation and lift the non-packed CP scope

Packed (THD) CP1-vs-CP8 at world64/EP64 trained 100 optimizer steps over
byte-identical topology-independent packs: loss-curve cosine 0.999999925,
mean absolute difference 0.00024, maximum 0.00069 - exceeding the
non-packed benchmark. Update the validated-configuration table and the CP
scope bullet accordingly.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(qwen3-8-flash-next): replace TileLang QSA with FlexAttention

Route-sparse GQA now executes through FlexAttention: selected token IDs are
scattered into a boolean membership table and evaluated via a BlockMask, so
one code path serves dense, packed (THD), and context-parallel layouts with
O(S) memory and no compiled-kernel dependency. The TileLang sparse-GQA
kernels (BSHD and THD) and their tests are removed; CPU keeps the gathered
PyTorch oracle. GDN gated RMSNorm and QK grouped RMSNorm fp32 chains gain
torch.compile variants on CUDA.

Validation on the 180B checkpoint (world64, EP64, 4k packed, GBS128):
- flex kernel: CP-slice vs full-row outputs bitwise identical; 0.18% rel L2
  vs the fp32 oracle
- 50-step CP1 vs CP8 packed SFT loss parity: mean |diff| 4.4e-4,
  max 1.3e-3, cosine 0.9999997 (matches the prior TileLang baseline)
- unit suite: 77 passed (CPU and H100)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* docs(qwen3-8-flash-next): use the official N-gram Embedding terminology

The Hugging Face model card and the DeepSeek Engram repository name this
component the N-gram Embedding; 'product-level embedding' was not a real
expansion of the checkpoint's ple identifier. Also align the coverage-page
summary table with the card's parameter and context-length figures and drop
stale fused-THD wording.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(qwen3-8-flash-next): align FlexAttention documentation

Remove a stale test import and make the packed-CP and FlexAttention
contracts match the implemented behavior, including the route-membership
workspace shape.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* docs(qwen3-8-flash-next): neutralize cluster-specific NCCL wording

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(qwen3-8-flash-next): align Engram projection input dtype

The model-owned table can retain fp32 master-weight storage while the block
computes in bf16; cast the lookup to the block activation dtype at the
projection boundary.

(cherry picked from commit ea975b221 and remapped from qwen4_exp)

Signed-off-by: Shuang Yu <shuangy@nvidia.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(qwen3-8-flash-next): align hyperconnection projection dtype

Cast the normalized residual stream to the projection weight dtype at the
two HC linear boundaries so an fp32 residual composes with bf16 projections
while the caller's residual dtype is preserved.

(cherry picked from commit e52f743b9 and remapped from qwen4_exp)

Signed-off-by: Shuang Yu <shuangy@nvidia.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* test(qwen3-8-flash-next): avoid pytest module collision

Rename the state-dict adapter test module so pytest does not confuse it
with the legacy bi-encoder test during full-suite collection.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(qwen3-8-flash-next): keep empty flex rows finite

Give route-empty padding queries a dummy kernel route before masking their
outputs so FlexAttention never evaluates a fully masked softmax row.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* docs(qwen3-8-flash-next): apply tech pubs copyedits

Polish the model overview, validation language, recipe sizing, and current
scope without changing the documented support contract.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* test(docs): map legacy Qwen4 architecture

Document that the checkpoint-facing Qwen4Exp architecture resolves to the
Qwen3.8-Flash-Next implementation covered by the existing model page.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

---------

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Signed-off-by: Shuang Yu <shuangy@nvidia.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Shuang Yu <shuangy@nvidia.com>