# 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)](https://github.com/fla-org/flash-linear-attention)
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`:

```bash
# 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:

```yaml
# 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:

```python
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:

```python
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
# 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:

```bash
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.