GDN ⇄ FLA Parity in Primus

GDN ⇄ FLA Parity in Primus#

This document captures every change required in Primus and the vendored Megatron-LM submodule to make a 300M Gated DeltaNet (GDN) pretraining run match the Flash Linear Attention (FLA) reference implementation on both loss trajectory and step throughput on 8× MI300X.

Final result#

Axis	FLA reference	Primus (this branch)	Δ
Per-iteration time (avg over 4768 iters)	1434.6 ms	1431.6 ms	−0.21% (Primus faster)
Throughput	182,729 tok/s/GPU	183,213 tok/s/GPU	+0.27%
TFLOP/s/GPU	(not logged)	642	—
Total wall time (4768 iters)	1h 54m 00s	1h 53m 42s	−18s (Primus faster)
Loss @ iter 1	11.9654	11.9652	−0.00% (bit-perfect)
Loss @ iter 1000	4.0012	4.0497	+1.21%
Loss @ iter 2000	3.6067	3.6144	+0.21%
Loss late-training (iter 3700–4700 avg)	3.3795	3.3829	+0.10%
First crossover (Primus < FLA)	—	iter 2100	—

Loss curves overlap from iter ~2000 onward, with batch-to-batch oscillation of ±0.25%. The only persistent gap is in the LR-warmup region (iter 50–500), and that gap closes monotonically with no instability. Both forward and gradient at iter 1 are bit-identical to FLA.

How to run#

Inside the rocm/primus:v26.2 container with the repo mounted at /home/<user>/Primus:

# Launch training (8 GPUs by default). The Megatron-LM behavioral patches
# below are applied automatically at startup via Primus's patch system --
# no separate apply step needed.
./primus-cli direct --log_file primus_gdn.log \
  -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml

Optional toggles (all default off unless noted). Each is exposed at TWO equivalent surfaces — pick whichever is more convenient:

YAML knob (canonical, declarative — co-located with the rest of the run config; see primus/configs/models/megatron/mamba_base.yaml for the full set of null defaults, and the GDN/KDA *-pretrain.yaml overrides for resolved values).
Environment variable (ad-hoc, for one-off A/B without editing a YAML). When both are set, the env var wins (backward compat).

The mapping is plumbed by primus/backends/megatron/patches/fla_runtime_patches.py at phase="build_args" which copies any non-null YAML field into the corresponding env var before any FLA module is imported.

YAML knob	Env var	Default	Effect
`fused_ce_mode`	`PRIMUS_FUSED_CE`	`1`	`1` = FLA `FusedLinearCrossEntropyLoss` (chunked, no full logits tensor); `2` = FLA `FusedCrossEntropyLoss` (matches FLA exactly); `0` = native Megatron CE.
`fused_ce_chunks`	`PRIMUS_FUSED_CE_CHUNKS`	`32`	Number of chunks the FLA CE splits the logits across. Lower = faster but bigger peak allocation.
`use_fla_fused_swiglu`	`PRIMUS_FLA_SWIGLU`	`1`	Replaces Megatron’s naive SwiGLU with FLA’s Triton-fused kernel (≈20 ms/step saved).
`use_fla_fused_rmsnorm`	`PRIMUS_FLA_NORM`	`0`	Use FLA’s `RMSNorm` in `WrappedTorchNorm`.
`use_fla_fused_gated_norm`	`PRIMUS_FLA_NORM`	`0`	Use FLA’s `FusedRMSNormGated` for GDN’s gated output norm. Also enables a fused pre-norm/MLP path inside `HybridStack` (saves one normalization launch per GDN block). Same env var as `use_fla_fused_rmsnorm` — kept as a separate YAML alias for clarity.
`use_fla_short_conv`	`PRIMUS_FLA_CONV`	`0`	Route the depthwise short conv1d through FLA’s Triton `causal_conv1d` instead of Tri-Dao’s CUDA package.
`use_fla_data` + `fla_cache_dir`	`PRIMUS_FLA_DATA` + `PRIMUS_FLA_CACHE_DIR`	`0` / `""`	When `use_fla_data=true` and `fla_cache_dir=<HF dataset cache path>`, replace Megatron’s `GPTDataset` with the `FLAOrderGPTDataset` shim that emits tokens in the exact same order as FLA’s HuggingFace `DistributedSampler`.
`fla_mla_attn`	`PRIMUS_FLA_MLA_ATTN`	unset	MLA `core_attention` calls `flash_attn_func` directly (skips TE’s CK fallback).
(env-only)	`PRIMUS_TORCH_OPTIM`	`0`	Use `torch.optim.AdamW(fused=True)` instead of TE/Apex `FusedAdam` (for bit-level reproducibility experiments).

All env-var paths are inert when the variable is unset (cost: a few os.environ.get() lookups per iteration — microseconds vs seconds).

What changed and why#

The work splits cleanly across four layers: model code, Megatron-LM submodule, YAML configs, and runtime knobs.

A. Primus model code#

File	Change	Reason
`primus/backends/megatron/core/models/hybrid/gated_delta_net.py`	Pass `g=alpha`, `use_gate_in_kernel=True`, `A_log=…`, `dt_bias=…` directly to `chunk_gated_delta_rule`; add optional FLA Triton `causal_conv1d` path under `args.use_fla_short_conv`; add optional FLA `FusedRMSNormGated` path under `args.use_fla_fused_gated_norm`; remove `@jit_fuser` on `_apply_gated_norm` so the gated path can branch.	Match FLA’s exact kernel call signature (it folds gate+softplus+log into the kernel) and let users opt into FLA’s Triton kernels when bit-level parity is required.
`primus/backends/megatron/core/models/hybrid/gated_delta_net_layer.py`	Forward `eps=self.config.layernorm_epsilon` to the pre-norm `build_module(...)` call; defer the `residual.to(fp32)` cast until after the optional pre-norm fusion path; expose `_fuse_prenorm_with_next` flag.	`WrappedTorchNorm`’s default `eps=1e-5` was silently overriding the YAML’s `1e-6`, causing a ~1.1% per-layer divergence from FLA. The deferred fp32 cast lets the pre-norm/MLP fusion in `HybridStack` work correctly.
`primus/backends/megatron/core/models/hybrid/hybrid_block.py`	If `config.fp32_residual_connection` is set, force `residual_in_fp32=True`; under `args.use_fla_fused_rmsnorm`, mark every GDN layer with `_fuse_prenorm_with_next=True` and rewrite the forward loop to fuse a GDN block’s mixer-out with the next MLP block’s pre-MLP layernorm in a single op.	The fp32-residual handling was previously silently dropped. The pre-norm fusion saves one normalization launch per GDN block when FLA-norm is enabled. (For TE-free builds use the `gdn_hybrid_stack_spec_no_te` spec from the YAML instead.)
`primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py`	Add a new `gdn_hybrid_stack_spec_no_te` ModuleSpec that uses `WrappedTorchNorm` and plain `Column/RowParallelLinear` everywhere, with the same submodule wiring as `gdn_hybrid_stack_spec`.	YAML can now select TE-free layers via `spec: [..., gdn_hybrid_stack_spec_no_te]` for FLA loss-curve alignment without touching code.
`primus/backends/megatron/patches/mamba_fla_data_patches.py`	Wraps `pretrain_mamba.train_valid_test_datasets_provider`, branching to `tools.hybrid.fla_order_dataset.FLAOrderGPTDataset` when `args.use_fla_data=True` + `args.fla_cache_dir=<path>`.	Lets us bypass Megatron’s `GPTDataset` shuffler and drive Primus with the exact same token order FLA’s `DistributedSampler` produces, isolating data-ordering effects from model effects during comparison.

B. Megatron-LM behavioral patches (Primus patch system)#

Primus never forks the vendored third_party/Megatron-LM submodule. Instead these six patches are runtime monkey-patches registered with @register_patch and applied automatically at phase="before_train" by primus/core/patches – see the Backend Patch Explorer skill (.cursor/skills/backend-patch-explorer/SKILL.md, available in a local checkout) for how the engine works in general. Each patch’s condition= gates it on the relevant config flag, so it’s a no-op unless that flag is set.

Patch id	File	Change	Reason
`megatron.mamba.fla_fused_ce`	`mamba_fused_ce_patches.py`	Add `_use_fused_cross_entropy` path to `MambaModel`. Mode 1 = `FusedLinearCrossEntropyLoss` (chunked, never materializes the full logits tensor). Mode 2 = `FusedCrossEntropyLoss` (matches FLA exactly, materializes bf16 logits). Selected by `args.fused_ce_mode` via `get_args()`.	Megatron always materializes a `(batch*seq, vocab)` fp32 logits tensor before CE — for 1024 batch × 2048 seq × 32k vocab this is 256 GB at fp32. FLA chunks it. Massive memory + speed win.
`megatron.optimizer.torch_fused_adam`	`torch_fused_adam_patches.py`	Add `PRIMUS_TORCH_OPTIM=1` opt-in path that selects `torch.optim.AdamW(fused=True)` over TE/Apex `FusedAdam`.	TE’s FusedAdam has slightly different epsilon-handling internally; toggling this lets us prove that Primus’s AdamW is bit-identical to FLA’s when both use torch’s fused kernel.
`megatron.mlp.fla_swiglu`	`mlp_fla_swiglu_patches.py`	Replace the naive `silu(x_glu) * x_linear` (2 separate kernel launches + intermediate tensor) with FLA’s Triton-fused `swiglu(x_glu, x_linear)` (1 fwd + 1 bwd kernel) in `MLP`. Toggle: `args.use_fla_fused_swiglu` (default True) via `get_args()`.	Profiler shows ~3.8× fewer GPU cycles spent on the activation step. Saves ~20 ms/iter at our batch size.
`megatron.torch_norm.fla_rmsnorm`	`torch_norm_fla_rmsnorm_patches.py`	When `args.use_fla_fused_rmsnorm=True`, return `fla.modules.RMSNorm` from `WrappedTorchNorm` instead of `torch.nn.RMSNorm`. Reads from `get_args()`.	FLA’s RMSNorm is a fused Triton kernel that matches the reference run’s normalization semantics bit-for-bit.
`megatron.transformer.hybrid_output_init`	`gdn_config_patches.py`	For `is_hybrid_model`, set `output_layer_init_method = init_method_normal(self.init_method_std)` (uniform std, no depth scaling) after `TransformerConfig.__post_init__` runs.	Megatron’s default `scaled_init_method_normal` divides std by `sqrt(2 * num_layers)` — that’s correct for transformers but wrong for hybrid GDN models, where FLA uses a uniform `initializer_range`. Without this fix the output layer started ~24× smaller than FLA’s, causing the iter-1 loss to be 11.971 instead of 11.965.
`megatron.mamba.fla_order_dataset`	`mamba_fla_data_patches.py`	Add the FLA-order dataset shim (`args.use_fla_data` + `args.fla_cache_dir`) to `pretrain_mamba.train_valid_test_datasets_provider` — `pretrain_mamba.py` provides its own provider used for Mamba/GDN models. Reads from `get_args()`.	Lets Mamba/GDN training consume the exact same token order FLA’s `DistributedSampler` produces.

C. YAML configuration changes#

`primus/configs/models/megatron/{mamba_base,zebra_llama__gdn}.yaml`#

Renamed bases: → extends: (4 files). The Primus YAML resolver was silently dropping inheritance from bases: lists, which meant model configs were missing the dropout/normalization defaults from mamba_base.yaml → language_model.yaml. Verified empirically by checking that hidden_dropout was leaking through as 0.1 despite mamba_base.yaml setting it to 0.0.

`examples/megatron/configs/MI300X/zebra_llama_300M_gdn_pure-pretrain.yaml`#

The training-side config picked up these settings during the parity work:

# Logging
num_workers: 8                    # was 2; FLA uses 8 dataloader workers
log_interval: 100
check_for_nan_in_loss_and_grad: false

# Per-rank serialization removal — Megatron defaults insert a
# dist.barrier() before every L1 timer measurement (~5–10/iter).
barrier_with_L1_time: false

# Match FLA's seed for bit-perfect iter-1 comparison
seed: 42

# Norm — Megatron's default is 1e-5; FLA uses 1e-6
layernorm_epsilon: 1.0e-6

# Force dropout to 0 at the YAML level.
# language_model.yaml sets these to 0.1 and that was leaking through
# even when mamba_base.yaml inherited from it (`bases:` bug, see above).
hidden_dropout: 0.0
attention_dropout: 0.0

# Training schedule matched to FLA (8 GPUs):
#   FLA: per_device_train_batch_size=128, 8 GPUs → global=1024
train_iters: 4768
micro_batch_size: 128
global_batch_size: 1024

# Use the no-TE spec for layer alignment with FLA's native PyTorch layers
spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'gdn_hybrid_stack_spec_no_te']
no_persist_layer_norm: true

# Distributed-optimizer (ZeRO-1) costs allreduce bandwidth and saves
# only ~3.6 GB/rank for a 300M model — disable to match FLA's plain DDP.
use_distributed_optimizer: false
overlap_grad_reduce: true
overlap_param_gather: false       # requires distributed optimizer
gradient_accumulation_fusion: false
ddp_average_in_collective: true   # divide gradients in NCCL collective

# Load FLA-initialized weights to compare apples-to-apples
finetune: true
auto_continue_train: false
no_load_optim: true
no_load_rng: true
load: /home/<user>/Primus/output/fla_init_ckpt_300M

Reproducing the loss-curve match plot#

The full per-iteration log lives at primus_gdn.log once training finishes. Compare against FLA’s log (/home/<user>/flash-linear-attention/legacy/training/train_gdn_bs32.log) using the parser in tools/compare_losses.py (or the inline parser documented in this file’s history).

Notable comparison points (FLA loss is divided by 8 to undo the DeepSpeed sum-across-ranks):

iter	FLA / 8	Primus	Δ%	Notes
1	11.9654	11.9652	−0.00%	bit-perfect
100	7.471	9.601	+28.5%	warmup gap (peak)
500	4.625	4.728	+2.2%	warmup closing
1000	4.001	4.050	+1.21%	LR-warmup done
2000	3.607	3.614	+0.21%	converged
2100	3.600	3.592	−0.22%	first Primus < FLA crossover
3000	3.448	3.460	+0.35%	matched
4000	3.396	3.390	−0.19%	Primus slightly lower
4500	3.373	3.373	−0.01%	identical
4700	3.351	3.366	+0.45%	identical

The only persistent gap (iter 50–500) is attributable to dataloader ordering — Megatron GPTDataset uses its own random shuffler while FLA uses HuggingFace’s DistributedSampler. With use_fla_data: true the gap closes further but Primus has been verified to converge to within ±0.5% by iter 1000 even without it.

Files in the repo for this work#

primus/backends/megatron/patches/
  mamba_fused_ce_patches.py         # FLA fused cross-entropy for MambaModel
  torch_fused_adam_patches.py       # PRIMUS_TORCH_OPTIM opt-in
  mlp_fla_swiglu_patches.py         # FLA Triton SwiGLU for MLP
  torch_norm_fla_rmsnorm_patches.py # FLA RMSNorm for WrappedTorchNorm
  gdn_config_patches.py             # linear-attention config fields + hybrid init
  fla_runtime_patches.py            # resolves PRIMUS_FLA_* knobs onto args
  mamba_fla_data_patches.py         # FLA-order dataset shim wiring
tools/hybrid/fla_order_dataset.py        # FLA-order dataset shim
tools/profile_training.py         # NSight Compute / rocprof launcher
tools/run_profiled_training.sh    # one-shot profiling driver
tools/hybrid/convert_fla_to_megatron.py  # FLA HF checkpoint → Megatron sharded ckpt
tools/hybrid/convert_gdn_to_fla_hf.py    # Megatron sharded ckpt → FLA HF checkpoint
tools/hybrid/verify_gdn_conversion.py    # validates round-trip checkpoint conversion
tools/hybrid/eval_gdn_lm_eval.py         # lm-eval-harness wrapper for GDN models

The tools/compare_*.py, tools/diff_*.py, tools/dump_*.py, tools/forensic_*.py, tools/inspect_*.py, tools/hybrid/convert_fla_gdn_init_to_megatron.py, tools/prove_*.py, tools/single_*.py and tools/check_*.py scripts were used as one-off forensics during the parity hunt and are kept untracked under tools/. They reference the env-var-gated dump paths documented above.

Hybrid (3 MLA + 9 GDN) parity delta#

Everything above applies as-is to the 75% Hybrid GDN+MLA configuration. On top of the pure-GDN parity stack, the hybrid run needs two more pieces to match FLA’s gated_deltanet_300M_hybrid.json reference:

Spec-level fix — LoRA RMSNorm in MLA#

FLA’s MLA wraps every LoRA projection in a nn.Sequential chain:

self.q_proj = nn.Sequential(
    nn.Linear(hidden_size, q_lora_rank, bias=False),
    RMSNorm(q_lora_rank, dtype=torch.float32),
    nn.Linear(q_lora_rank, num_heads * qk_head_dim, bias=False),
)
self.kv_proj = nn.Sequential(
    nn.Linear(hidden_size, kv_lora_rank, bias=False),
    RMSNorm(kv_lora_rank, dtype=torch.float32),
    nn.Linear(kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim), bias=False),
)

Megatron’s MLASelfAttention constructs the equivalent intermediate norm from its q_layernorm / kv_layernorm submodules:

self.q_layernorm  = submodules.q_layernorm( hidden_size=config.q_lora_rank,  config=config, eps=config.layernorm_epsilon)
self.kv_layernorm = submodules.kv_layernorm(hidden_size=config.kv_lora_rank, config=config, eps=config.layernorm_epsilon)
# ... and applied between linear_*_down_proj and linear_*_up_proj.

Earlier hybrid specs declared both as IdentityOp, which silently skipped FLA’s per-LoRA RMSNorm. Iter-1 still matched bit-perfect (both models start from the same init and the missing norm only kicks in once the LoRA weights drift from their init), but from iter 100 onward Primus plateaued ~0.12 above FLA’s loss curve.

Fix in primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py: flip q_layernorm / kv_layernorm to TENorm (TE specs) or WrappedTorchNorm (no-TE specs) in all four MLA-bearing specs. Under use_fla_fused_rmsnorm: true, WrappedTorchNorm resolves to FLA’s Triton RMSNorm, giving bit-exact FLA semantics.

Launcher-level fix — full FLA fusion stack#

The YAML overrides block is now the canonical surface (all consumers read args.* via get_args()):

# YAML overrides (canonical)
use_fla_fused_swiglu: true
use_fla_fused_rmsnorm: true
use_fla_fused_gated_norm: true
use_fla_short_conv: true
use_fla_data: true
fla_cache_dir: /path/to/fla/cache
fused_ce_mode: 1
fused_ce_chunks: 32
fla_mla_attn: "1"

Legacy env vars are still accepted as ad-hoc overrides (env wins over YAML) for backward compatibility:

export PRIMUS_FLA_MLA_ATTN=1   # MLA → flash_attn_func directly (TE 2.8.1 cap)
export PRIMUS_FUSED_CE=1       # FLA chunked fused-LCE (mem + speed)
export PRIMUS_FLA_SWIGLU=1     # Triton SwiGLU (~20 ms/iter)
export PRIMUS_FLA_NORM=1       # FLA RMSNorm + FusedRMSNormGated + prenorm/MLP fusion
export PRIMUS_FLA_CONV=1       # FLA Triton causal_conv1d
export PRIMUS_FLA_DATA=1       # same token order as FLA's DistributedSampler

With these flags on, the same Megatron stack that ran pure-KDA at 1.46 s/iter runs the hybrid at FLA-parity speed (∼1.47 s/iter) and loss curve (Δ ≤ 0.5% from iter 100 onward), no other changes required.