From 1940aa7b96cdfcdb84de25b98b12a97f8fec1c1b Mon Sep 17 00:00:00 2001 From: echobt <154886644+echobt@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:38:00 +0000 Subject: [PATCH 1/2] docs: sync recipe 2.1 and LoopMoE example Mirror BASE v2.1 miner contract (G2 leaf, 1B/4-GPU, dual cap) and add the LoopMoE AutoModel example. No control-plane source. --- README.md | 20 +- docs/README.md | 5 +- docs/getting-started.md | 13 +- docs/prism.md | 387 ++++- docs/scoring.md | 12 +- examples/loopmoe/README.md | 37 + examples/loopmoe/__init__.py | 4 + examples/loopmoe/automodel.base | 1 + examples/loopmoe/automodel.patch | 2333 +++++++++++++++++++++++++++++ examples/loopmoe/ddp_worker.py | 47 + examples/loopmoe/entry.py | 846 +++++++++++ examples/loopmoe/kernels.py | 867 +++++++++++ examples/loopmoe/model.py | 521 +++++++ examples/loopmoe/prism.toml | 1 + examples/loopmoe/requirements.txt | 12 + 15 files changed, 5042 insertions(+), 64 deletions(-) create mode 100644 examples/loopmoe/README.md create mode 100644 examples/loopmoe/__init__.py create mode 100644 examples/loopmoe/automodel.base create mode 100644 examples/loopmoe/automodel.patch create mode 100644 examples/loopmoe/ddp_worker.py create mode 100644 examples/loopmoe/entry.py create mode 100644 examples/loopmoe/kernels.py create mode 100644 examples/loopmoe/model.py create mode 100644 examples/loopmoe/prism.toml create mode 100644 examples/loopmoe/requirements.txt diff --git a/README.md b/README.md index 086cb5e..9510473 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,10 @@ PRISM is a research challenge on a pinned [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel) base: you fork the operator pin, edit under that tree, and submit a **unified git diff**. The operator applies your patch fail-closed, then re-executes training on a -miner-funded Lium GPU pod against a pinned FineWeb-Edu shard. Score is pure -**bits-per-byte** (bpb, lower is better). There is **no** miner Docker image, -no CVM, no on-chain write from miners — HTTP submit only. +miner-funded **4-GPU** Lium pod against a pinned FineWeb-Edu shard. Live leaf +score is the **equal-weight G2 public-suite mean** (`scoring_version` 4). +There is **no** miner Docker image, no CVM, no on-chain write from miners — +HTTP submit only. | | | |---|---| @@ -35,7 +36,7 @@ no CVM, no on-chain write from miners — HTTP submit only. | Production gateway | `https://chain.joinbase.ai` | | Staging gateway | `http://staging.api.joinbase.ai` | | Submit path | `/challenge/prism/v1/submissions` | -| Recipe | **2.0.0** — AutoModel pin + patch (`automodel@v0.5.0`) | +| Recipe | **2.1.0** — AutoModel pin + patch (`automodel@v0.5.0`), 4-GPU CUDA 13/TE | | Live GPU | Miner-funded Lium — pass `X-Lium-Api-Key` | This repository holds **miner documentation and examples only**. Control-plane @@ -46,9 +47,10 @@ source lives in [BaseIntelligence/base](https://github.com/BaseIntelligence/base 1. Read [Getting started](docs/getting-started.md) (or the [full guide](docs/prism.md)). 2. `GET /v1/recipe` — copy `automodel_pin_id`, `automodel_git_commit`, and caps. 3. Checkout that AutoModel commit → edit → `git diff > automodel.patch`. -4. Pack `automodel.base` + `automodel.patch` (+ optional `prism.toml`) and submit - with your hotkey + **`X-Lium-Api-Key`** — see [Submit](docs/submit.md). -5. Poll events until `terminated`, then check your bpb — see [API](docs/api.md). +4. Pack `automodel.base` + `automodel.patch` (+ optional `prism.toml` / + `requirements.txt`) and submit with your hotkey + **`X-Lium-Api-Key`** — + see [Submit](docs/submit.md). Worked example: [LoopMoE](examples/loopmoe/). +5. Poll events until `terminated`, then check G2 benches — see [API](docs/api.md). ```bash export GATEWAY=https://chain.joinbase.ai @@ -68,8 +70,8 @@ curl -sS -X POST "$GATEWAY/challenge/prism/v1/submissions" \ ## The three things miners get wrong 1. **Legacy 1.x ZIPs** — `architecture.py` + `training.py` (or training-only - `arch_id`) return `400 unsupported_layout` / `recipe_version` on live 2.0. - Ship `automodel.base` + `automodel.patch` only. + `arch_id`) return `400 unsupported_layout` / `recipe_version` on live 2.1. + Ship `automodel.base` + `automodel.patch` only. 2. **Wrong pin / stale diff** — `automodel.base` must equal live `automodel_pin_id` (`automodel@v0.5.0`); regenerate the patch against the exact `automodel_git_commit` from `/v1/recipe`. diff --git a/docs/README.md b/docs/README.md index 21aa4d6..f05ab92 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,7 +1,10 @@ # PRISM miner docs -Live recipe is **2.0.0**: submit an AutoModel pin id + unified git diff +Live recipe is **2.1.0**: submit an AutoModel pin id + unified git diff (`automodel.base` + `automodel.patch`), not a free-form two-script ZIP. +Pods expose **4 GPUs** (CUDA 13 / Transformer Engine). Caps: **≤ 1B params**, +**3.0e18 attested FLOPs** + 5 h wall. Leaf scoring is **G2 benchmarks** +(`scoring_version` 4) unless operators flip a documented knob. | Page | What it covers | |------|----------------| diff --git a/docs/getting-started.md b/docs/getting-started.md index cb3e011..20d2054 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -1,9 +1,9 @@ # Getting started -## The contract (recipe v2.0.0) +## The contract (recipe v2.1.0) You do **not** ship a free-form `architecture.py` / `training.py` project. -Live recipe **2.0.0** accepts only a pin id plus your unified diff against that +Live recipe **2.1.0** accepts only a pin id plus your unified diff against that pin: ```text @@ -25,9 +25,12 @@ prism.toml # optional — entry / model-config knobs 5. Write `automodel.base` as a single line equal to `automodel_pin_id`, pack the ZIP, and `POST /v1/submissions` with your hotkey + **`X-Lium-Api-Key`**. -Models must stay **≤ 350M parameters**. The pod has **no network** -(`unshare --net`) beyond the operator-owned dataset pull — do not call Hub -downloads from miner code. +Models must stay **≤ 1B parameters**. Recipe-v10 pods expose four GPUs; +train from `ctx["train_stream"]` (rank 0 owns the stream under DDP). Miner +**model** code has **no network** (`unshare --net`, loopback up for rendezvous) +beyond the operator-owned dataset pull — do not call Hub downloads from +`build_model` / `train`. You may ship `requirements.txt` or `pyproject.toml` +for a network-on install phase before that sandbox. **Legacy recipe 1.x is rejected on live.** Two-script ZIPs (`architecture.py` + `training.py`), 1.3 source-tree ZIPs, and training-only diff --git a/docs/prism.md b/docs/prism.md index a0a0818..680679c 100644 --- a/docs/prism.md +++ b/docs/prism.md @@ -3,18 +3,16 @@ # Prism challenge — HTTP AutoModel patch submit **challenge_id:** `prism` -**scoring_version:** `2` live (bpb-only; LLM review is an anti-cheat gate, not a grader). **v3 (opt-in, shadow-by-default):** composite scoring runs alongside — your run is also measured on the G1–G8 battery; see *v3 scoring* below. -**recipe_version:** `2.0.0` (pinned [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel) base + miner unified diff; legacy 1.x layouts rejected on live) +**scoring_version:** `4` live (equal-weight G2 public-suite accuracies → lattice; LLM review is an anti-cheat gate, not a grader). **v3 harness (default):** every scored run executes the **G1–G8 battery**; the leaf uses G2 benches while `PRISM_SCORING_MODE=benchmarks` (default). Legacy `shadow` = bits/token bpb; `composite` = full G1–G8 lattice when anchors are ready. +**recipe_version:** `2.1.0` (pinned [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel) diff + 4-GPU CUDA 13/TE pod + attested dual cap; legacy 1.x layouts rejected) **Path:** HTTP only — **no Phala/CVM** -Normative docs (BASE monorepo): -[`PRISM.md`](https://github.com/BaseIntelligence/base/blob/main/docs/PRISM.md), -[`PRISM_RECIPE.md`](https://github.com/BaseIntelligence/base/blob/main/docs/PRISM_RECIPE.md). +Normative docs (BASE monorepo): [`PRISM.md`](https://github.com/BaseIntelligence/base/blob/prism-v2.1-scoring/docs/PRISM.md), recipe [`PRISM_RECIPE.md`](https://github.com/BaseIntelligence/base/blob/prism-v2.1-scoring/docs/PRISM_RECIPE.md). ## What you submit A **ZIP** (preferred) — or JSON with the same members / `zip_base64` — that is -**not** a free-form `architecture.py` / `training.py` project. Recipe **2.0.0** +**not** a free-form `architecture.py` / `training.py` project. Recipe **2.1.0** accepts only an AutoModel pin id plus your git diff against that pin: ```text @@ -36,9 +34,59 @@ prism.toml # optional — entry / model-config knobs 5. Write `automodel.base` as a single line equal to `automodel_pin_id`, pack the ZIP, and `POST /v1/submissions` with your hotkey + **`X-Lium-Api-Key`**. -Models must stay **≤ 350M parameters**. The pod has **no network** -(`unshare --net`) beyond the operator-owned dataset pull — do not call Hub -downloads from miner code. +Models must stay **≤ 1B parameters**. Miner **model code** (build/train/ +eval) runs with **no network** (`unshare --net`) beyond the operator-owned +dataset pull — do not call Hub downloads from `build_model` / `train`. + +**Bring your own dependencies (recipe-v10).** The pod image is a complete +CUDA 13 base with PyTorch, a build toolchain (`nvcc`, `ninja`), Transformer +Engine (NVFP4 training), and common accelerators. You may additionally ship +**one** of: + +- `requirements.txt` — installed with `pip install -r requirements.txt` +- `pyproject.toml` — installed with `pip install .` + +by **adding the file at the repo root in your `automodel.patch`** (or at +the ZIP root on the legacy two-script path). It is installed in a +**network-on install phase before** your model code is sandboxed — so +`flash-attn`, `mamba-ssm`, custom Triton/CUDA kernels, etc. compile and +install, then train/eval run offline. `requirements.txt` wins if you ship +both. Check `GET /v1/recipe` for `pod_image_ref`, `miner_install_supported`, +`miner_deps_members`, and `install_timeout_secs`. + +**Four-GPU train contract.** Recipe-v10 pods expose four RTX 5090 GPUs by +default through `ctx["gpu_count"]`; the organizer eval stays on GPU 0. +Training must consume global batches from the harness-owned +`ctx["train_stream"]`, because that stream enforces step/wall/FLOPs caps and +owns token/byte accounting for G6. For DDP, keep rank 0 as the stream owner +and scatter/shard each accounted global batch to workers over the local +process group. Do not let each worker create an independent dataset stream: +v3 rejects a trainer that returns with zero harness-accounted tokens. The +isolated network namespace brings `127.0.0.1` loopback up for rendezvous but +has no external route. + +**Resubmit at will on your own failures.** If your dependency install fails +(`install_deps`) or your `training.py` crashes at build/train time +(`train_script`), the run fails **without burning your one-submission +slot** — fix the manifest or the script and resubmit immediately, no time +window. (Operator-side infra hiccups keep the existing 30-minute resubmit +window.) + +**Bring your own tokenizer — verified.** The tokenizer is part of your +submission: ship `tokenizer/` files in your tree (≤ 12 files, ≤ 8 MiB, +loaded offline with `AutoTokenizer.from_pretrained(dir, +local_files_only=True)`) or export `def build_tokenizer(ctx)` next to +`build_model` in `architecture.py` (train/wrap anything, offline). The +harness hands it back as `ctx["tokenizer"]` / `ctx["vocab_size"]`; the +`gpt2` pin is only the default when you ship nothing. Two things keep this +fair: (1) G1 scores **bits/byte**, tokenizer-neutral — exotic vocabularies +buy you nothing on the headline metric; (2) every run emits an objective +**tokenizer card** (compression on a fixed probe, roundtrip fidelity, +vocab-shape scan) that the LLM anti-cheat review reads. A tokenizer +engineered to game metrics — multi-word answer phrases as single tokens, +vocab stuffed with eval-looking strings, a `decode()` that rewrites output, +memorizing compression — is a **cheat** (`tokenizer_gaming`, Score 0). A +merely weak tokenizer is not a cheat; it just hurts your own score. **Legacy recipe 1.x rejected on live.** Two-script ZIPs (`architecture.py` + `training.py`), 1.3 source-tree ZIPs, and training-only @@ -51,6 +99,13 @@ frameworks. bypass those hooks fail review (`missing_telemetry_hooks`, zero score, terminal). +**Example: LoopMoE (4-GPU DDP + NVFP4).** A reference AutoModel patch that +honors `ctx["train_stream"]`, rank-0 stream ownership, and optional +Transformer Engine NVFP4 lives at +[`examples/loopmoe/`](../examples/loopmoe/). Pack `automodel.base` + +`automodel.patch` (+ optional `prism.toml` / `requirements.txt`) as in +this document. It is an example, not a scored baseline. + **Diff visibility.** After intake, inspect your applied delta at `GET /v1/submissions/{id}/diff` (full unified diff + diffstat / classification). @@ -68,16 +123,29 @@ X-Lium-Api-Key: ``` The key is held in master memory for that submission and may also land in a -**short-TTL encrypted seal file** on the master host (never in Postgres, never -logged) so a control-plane restart can still stop your pod. Missing key on -live → `400 missing_lium_api_key`. Cost guardrails (`max_price_per_hour`, -lifetime) still apply so a bad key cannot rent unbounded SKUs through the -orchestrator. - -If the challenge process restarts mid-run, your submission is marked failed -promptly with `control_plane_restart` / `harness_detached` (not left `running` -for hours). When the seal is missing, stop the pod yourself on Lium, then -resubmit with `X-Lium-Api-Key`. Poll `GET /v1/submissions/{id}/events` and +**TTL-bounded encrypted seal file** on the master host (default ≥36h; never in +Postgres, never logged). Master **re-seals** on measure start and heartbeats +so a full 6h train wall cannot outlive the seal across a control-plane +restart. Missing key on live → `400 missing_lium_api_key`. Cost guardrails +(`max_price_per_hour`, lifetime) still apply so a bad key cannot rent +unbounded SKUs through the orchestrator. + +**Pod lifetime ceiling is 8.5h** (was 7h). You are billed for time actually +used, not the ceiling, so a run that finishes early costs the same as +before — the raise exists because the old 7h could **terminate a +full-budget submission mid-eval** and lose the whole rental. The ceiling has +to contain build (≤15m) + your 6h train wall + checkpoint (≤30m) + the eval +phase (≤1.5h) ≈ 8.3h. The eval battery itself runs under one global 1h +budget with per-group shares, and if a group hits its ceiling the run +reports it (`budget.truncated` / `budget.partial_groups` in the battery +blob) rather than silently scoring fewer items. + +If the challenge process restarts mid-run while your Lium pod is still +training/evaling, master **reattaches** quietly (same submission id; pod is +not killed). You only see `control_plane_restart` / `harness_detached` when +the pod is already dead or the sealed key cannot be restored and master +cannot talk to Lium — then stop the pod yourself and resubmit with +`X-Lium-Api-Key`. Poll `GET /v1/submissions/{id}/events` and `GET /v1/submissions/{id}/logs?since=` for live stage heartbeats and harness tails while the run is healthy. @@ -110,12 +178,19 @@ Inspect recipe + AutoModel pin before coding: curl -sS "$BASE_GATEWAY/challenge/prism/v1/recipe" ``` -Live recipe **2.0.0** advertises `version: "2.0.0"` and AutoModel pin fields +Live recipe **2.1.0** advertises `version: "2.1.0"` and AutoModel pin fields (`automodel_pin_id` = `automodel@v0.5.0`, `automodel_repo_url`, `automodel_git_ref`, `automodel_git_commit`, `automodel_content_sha256`), -plus caps such as `train_hours_cap: 6.0`, `max_train_steps: 20000`, -`max_params: 350000000`, FineWeb dataset pin, and `pin_hex` (sha over the -versioned descriptor). Trust `/v1/recipe`, not marketing chart labels. +plus caps such as `train_flops_cap: 3.0e18` (the budget currency), +`train_hours_cap: 5.0` (anti-DoS wall), `min_spend_fraction: 0.5` +(voluntary-stop floor; a step/wall/FLOPs-bound run stays eligible), +`max_train_steps: 20000`, `max_params: 1000000000`, FineWeb dataset pin, +and `pin_hex` (sha over the versioned descriptor). Trust `/v1/recipe`, +not marketing chart labels. + +The FLOPs probe retries OOM at progressively smaller row counts down to one +sequence and reports whether it reduced the batch; a deliberately +memory-heavy model cannot turn probe OOM into an unmetered train. `POST /v1/submissions` is idempotent by `submission_id` (hash of **pin id ‖ `0x00` ‖ patch bytes**). @@ -133,10 +208,24 @@ versioned descriptor). Trust `/v1/recipe`, not marketing chart labels. - If your hotkey **leaves the metagraph**, the watcher reopens your slot(s) automatically — resubmit under your new uid. - Infra failures (Lium pod, review/similarity/LLM infra) **auto-retry up to 3 - times**; cheat / rejected verdicts are terminal. After an infra failure - (`ChallengeInternal`), you may **resubmit within 30 minutes** (new POST or - `POST /v1/submissions/{id}/retry`). After 30 minutes the slot stays blocked - until your hotkey leaves the metagraph. + times**; harness `EVAL_FAIL` (miner/model code) is terminal for that attempt + and is **not** auto-retried. Cheat / rejected verdicts are terminal. After an + infra failure (`ChallengeInternal`), you may **recover within 30 minutes** + via `POST /v1/submissions/{id}/retry` with **`X-Lium-Api-Key`** (required on + live when another GPU rent is needed). After 30 minutes the slot stays + blocked until your hotkey leaves the metagraph. + +### Retry vs re-POST + +| Action | When | Headers | +|--------|------|---------| +| Re-POST the **same** ZIP | Always safe | Same as submit | Returns `200 already-queued` — **no new GPU run**; does not recover a failed row | +| `POST /v1/submissions/{id}/retry` | Row status is **`failed`** only | **`X-Lium-Api-Key`** on live (infra recovery); admin Bearer for operator non-infra retries | Requeues measure; wrong/missing Lium key → `400 missing_lium_api_key` | +| `/retry` on non-failed | — | — | `409 not_failed` — hotkey or Bearer alone does not change that | + +Do **not** expect `X-Miner-Hotkey` or admin Bearer alone to fund a new Lium +pod. Seal TTL is ≥36h and master re-seals on measure + heartbeats; the key is +kept across measure Err so auto-/miner-retry can re-rent without a new submit. ## Anti-copy rule (patch / delta) @@ -186,8 +275,11 @@ submission and never rents a Lium pod. ## Scoring (summary) -Final leaf score is pure bits-per-byte (bpb) on the lattice `[0, SCORE_MAX]`. -The shared **agentic** gate (AST + metrics/receipt) hard-zeros `cheat` / +Final leaf score (live `scoring_version` **4**) is the **equal-weight mean of +available G2 public accuracies** mapped to `round(SCORE_MAX × mean)` — not +bits/token bpb. Tokenizer length cannot farm the rank. Bits/token bpb and +tokenizer-neutral `bits_per_byte` remain recorded for display / G1. The shared +**agentic** gate (AST + metrics/receipt) hard-zeros `cheat` / `suspicious`. Cheap LLM similarity hard-zeros `Copied`, and `Suspicious` only when confidence `≥ 0.9` with non-generic evidence (below that — e.g. 0.7 citing RMSNorm/SwiGLU/LayerNorm — does **not** wipe your score). Copy/similarity @@ -199,20 +291,135 @@ quality is coherence-only, not a grader. Public gallery/leaderboard show champions only. **Competition (temporary):** emission uses **your own best training score only** — architecture-owner credit (rewarding arch owners when others train -well on their code) is **disabled** for now so the best-BPB trainer keeps +well on their code) is **disabled** for now so the best-scoring trainer keeps Prism's weights. Emission remains **winner-take-all**: only the single highest own score that epoch receives Prism's share (50% of the subnet); ties break by -lexicographically smallest hotkey. Scores first land in the leaf +lexicographically smallest hotkey. Two **v2.1 opt-in** emission knobs exist +but are **off by default** (operators announce any flip): `top3` mode pays +the top three positive scores at 100 % / 50 % / 25 % of their own lattice +score instead of winner-take-all, and an architecture-owner split can carve +up to 50 % of the winner's leaf to the **registry owner** of the winning +architecture — publishing a strong architecture that someone else trains to +the top then earns you a share. A third mode, `sig` (**significance-gated**, +also off by default), is described in +[what actually earns emission](#what-actually-earns-emission-significance-gating) +below. Scores first land in the leaf set emitted at the first chain-epoch boundary **after** your run finalizes (a long train that crosses epochs is normal — outbox assignment is exactly once). Positive scores then keep participating in later epochs' competition sets until a better valid score supersedes them (WTA still collapses to one leaf winner). -The global-best model (sources + `ARTIFACT.json` / checkpoint release) is -published to +The global-best model by **G2 lattice score** (sources + `ARTIFACT.json` / +checkpoint release) is published to [`BaseIntelligence/prism`](https://github.com/BaseIntelligence/prism) -`top-model/`. - -## v3 scoring (shadow-by-default) +`top-model/` and (when configured) a HuggingFace model repo +`BaseIntelligence/top-prism-architecture` (custom-arch / AutoModel novelty + +weights, `trust_remote_code`). See [`PRISM.md`](https://github.com/BaseIntelligence/base/blob/prism-v2.1-scoring/docs/PRISM.md). + +## What Prism does and does not claim about your architecture + +Read this before optimizing anything, because it tells you what the ranking +means. + +**Prism ranks architectures at a pinned, small budget.** Every submission trains +under the same fixed recipe, the same fixed data, the same wall-clock cap, on the +same GPU class. That is a genuine, well-controlled comparison — the pinned-recipe +discipline is what makes it meaningful at all — and it is the thing the leaderboard +measures. + +**Prism does not claim to select architectures that will scale.** This is not +modesty; it is a measured result we would rather publish than hide. Tay et al., +*Scaling Laws vs Model Architectures* (EMNLP Findings 2023, +[arXiv 2207.10551](https://arxiv.org/abs/2207.10551)) pretrained **>100 models +across 10 architectures from 15 M to 40 B parameters** and found: + +- "**The best performing model can fluctuate at different scales.**" +- The **vanilla Transformer has the best scaling exponent** while *not* being the + best at every individual compute point — the winner at one budget is not + necessarily the best scaler. +- **Concrete rank flips:** Evolved Transformer beats vanilla at small scale on + downstream tasks and falls behind when scaled up; **ALBERT scales negatively + downstream** (α = −0.12), and ALBERT's mechanism is cross-layer weight sharing, + the same family as looped / recurrent-depth designs. + +So a win here is evidence that your architecture is better **at this budget**, not +a prediction about 70 B. The authors also state the converse, which is Prism's +honest positive claim: not every practitioner needs models that scale to billions, +and inductive biases tailored to small or low-compute regimes are valuable in +their own right. That is the regime Prism measures. + +Two practical consequences for you: + +- Improvements that only appear at larger scale will not be visible here, and + that is a limitation of the instrument, not a judgement on your idea. +- Tuning tricks that exploit the pinned budget specifically may win here without + transferring. We would rather you know that than discover it later. + +## What actually earns emission (significance gating) + +Live emission today is **winner-take-all**: the single highest own score takes +Prism's share. The `sig` mode described here is **implemented but off by +default** — operators announce any flip, and it cannot be enabled until the +run-to-run (seed) noise floor has been measured and published. It is documented +now so you can see where the incentives are going. + +Under `sig`, **beating the champion requires clearing a measurement-uncertainty +margin, not just posting a better point estimate.** + +**Why.** Two runs of the *same* architecture do not produce the same number. +Eval sampling and training-seed noise both move the score. So a challenger that +scores 0.1 % better has not shown it is better — it may simply have drawn a +luckier run. Under the old rule that coin flip won the entire share, which is +also exactly what makes copying the champion profitable: a copy has the *same* +true quality, so it wins the flip about half the time. Requiring real evidence of +improvement removes that. + +**How the comparison works, in plain terms.** + +1. You and the champion are scored on the **same** private eval slice. +2. Your scores are compared **example by example**, not as two totals. Hard + examples are hard for both models, and pairing cancels that out. +3. Differences smaller than a fixed **dead zone** (0.01 in absolute metric units + — bits/byte, or accuracy) are treated as *undecided*. Hairline differences do + not vote. +4. You must win **≥ 55 % of the decided examples** — and not on the point + estimate: on a **99 % lower confidence bound** from a 10 000-resample + bootstrap with a fixed, published seed. Anyone can recompute the verdict. +5. Your **average margin** on decided examples must also clear the dead zone, so + you cannot win a majority of near-ties while being much worse where you lose. +6. There must be at least **100 decided examples**. Below that the win rate + cannot be estimated closely enough to mean anything, and the champion holds. + On a task where everyone scores nearly the same, that is the normal outcome: + the comparison refuses rather than crowning a coin flip. + +One consequence worth knowing: you and the champion must have been measured on +the **same** eval slice for any of this to run. If the slice rotated between the +champion's run and yours, there is no valid comparison and the champion holds +until it is re-measured on your slice. That re-measurement is the operator's job, +not yours. + +The bar is 55 % rather than something higher on purpose: a genuinely better +architecture with a wide per-example spread sits near 55 %, so demanding much +more would select for **low-variance submissions instead of good ones**. + +**What you get paid.** Prism's share splits: **60 %** champion (dropping toward +a 50 % floor if the win is real but marginal, with the difference **burned**), +**15 / 10 / 5 %** to ranks 2–4, and **10 %** split across up to five entries that +pass every gate and **hold the best measured value on any single axis** `g1..g8`. + +That last pool is the one worth understanding. You do **not** have to win overall +to earn from it. If you are third on the composite but **first on G3** +(associative recall) or **first on G7** (inference cost), you produced real +information and you are paid for it. Note what this is not: there is **no reward +for being different**. Renaming variables, reordering statements, or otherwise +looking novel earns nothing — the axes are real measurements, and you have to +actually be best on one. Anything unallocated **burns** rather than being +redistributed. + +Also under `sig`: the champion is **re-measured** on fresh private slices at +unannounced times (operator-funded, eval-only). If a champion's score was propped +up by fitting the public anchors, that shows up and costs it the title. + +## v3 scoring (battery always; leaf mode via env) Recipe ≥ 1.3.0 harnesses run a **two-phase pod flow**: your code trains (`phase=train`), checkpoints, and only then does the operator stage private @@ -224,6 +431,16 @@ inference efficiency (G7), and training stability/µP (G8). Everything the battery reports is organizer-measured (**Zone A**, `org.*`) and is computed inside the harness — your code never emits it. +The v3 metric surface is structurally complete: G1 includes code, prose, +math, fresh crawl, and key-token bits/byte; G6 includes byte-denominated AUC, +bytes-to-threshold, and bpb at half of the organizer FLOPs cap; G7 includes +measured-or-censored 32k TTFT/TPOT/state, board energy, throughput, and +reasoning throughput; G8 emits loss stability and µP LR-transfer. Unsupported +32k/OOM/power cases receive explicit worst-case censored values rather than +silently disappearing. Live emission nevertheless remains G2 benchmarks + +WTA until operators announce calibrated v3 anchors and a separate governance +flip. + **G5 is pretrain-only (recipe ≥ 1.4.0).** The long-context group scores a **base LM**, not an instruction-tuned chat model: completion-style / few-shot base prompts, short exact-match or multiple-choice logprob — @@ -252,17 +469,99 @@ terminal-loss band) and the cross-miner cohort, and land a stored verdict (`ok` / `flagged` / `quarantined`) — verdicts are evidence, never an auto-zero. Malformed or over-cap envelopes reject `422` and store nothing. -While `PRISM_SCORING_MODE=shadow` (default) the leaf score stays pure bpb, -bit-identical to v2. After the reference baselines are measured and the -anchor set is pre-registered, governance may flip to `composite`: group -scores are anchor-normalized, gate-filtered (`g3 ≥ 0.25`, `g8 ≥ 0.5`, -budget + CI gates), combined as a weighted geometric mean, and ranked by -the bootstrap lower-confidence bound +While `PRISM_SCORING_MODE=benchmarks` (default) the leaf score is the +**equal-weight mean of available G2 public accuracies** (HellaSwag, ARC-E/C, +PIQA, WinoGrande, BoolQ, LAMBADA strict when present, OpenBookQA), mapped to +`round(SCORE_MAX × mean)`. Missing every listed bench → `0` (fail-closed). +**Tokenizer length no longer farms the rank** — bits/token bpb is still +recorded (and tokenizer-neutral `bits_per_byte` feeds G1) but does **not** +drive emission. `PRISM_SCORING_MODE=shadow` restores legacy pure bits/token +bpb (v2). After the reference baselines (**Transformer++** and +**hybrid delta** — published in-repo under `crates/prism-recipe/baselines/`) +are measured and the anchor set is pre-registered, governance may flip to +`composite`: group scores are anchor-normalized (**arithmetic** mean within +each group; a single zero sub-metric does not zero the group), gate-filtered +(`g3 ≥ 0.25` currently **disarmed** until G3 item counts stabilize; +`g8 ≥ 0.5`, budget + CI gates), combined as a weighted +**geometric** mean across groups (`C = ∏ g_k^{w_k}` — a full group score of 0 +collapses C), and ranked by the bootstrap lower-confidence bound (`lattice = round(SCORE_MAX × max(0, C − 1.645·SE))`). Inspect the anchor registry and pre-registration commits at `GET /v1/anchors` and `GET /v1/preregistration`; per-run Zone A / Zone B rows at `GET /v1/submissions/{id}/metrics?zone=a|b`. +**G8 µP probe.** The stability sweep builds 1× and 4× width from a **fixed +small** width/depth base (not your full ≤1B scored model), then scales with +`ctx["prism_width_multiplier"]`. Honor top-level / `arch` geometry overrides +and that multiplier in `build_model` (reference baselines do) or the sweep +fail-closes `org.g8.mup_lr_stability = 0.0`. + +### v2.1 battery additions (anchor set v1, opt-in) + +Two extra organizer-measured keys ship with the v2.1 harness on every real +run (inert until operators select anchor set v1; you will see them in +`GET /v1/submissions/{id}/metrics?zone=a`): + +- `org.g7.reasoning_throughput` — mean G4 accuracy × decode toks/s. + Compute-normalized reasoning: architectures that spend extra inference + compute to reason (loops, adaptive depth, recursion) are credited for + the accuracy they buy in the same key that charges its cost — raw + throughput alone no longer structurally penalizes them. +- `org.g8.mup_scaling_slope` — a local scaling-exponent probe measured on + the existing µP 1×/4× width sweep (how fast your architecture improves + with scale). Support the `prism_width_multiplier` build knob (already + required for G8) and this costs you nothing extra; a failed sweep + fail-closes the key to 0.0. + +Practical consequence for architecture design: under anchor set v1, +"thinks more when it's hard" designs and "scales steeper" designs earn +score on dedicated axes instead of only paying G7/G6 penalties. + +### v2.2: LAMBADA scored strict (anchor set v2, opt-in) + +The G2 LAMBADA item used to be a 4-way multiple choice against random +distractor words — nearly free points (0.95+ for everyone, 0.985 for the +GPT-2 Large reference), because LAMBADA's gold word is uniquely determined +by its long context. The harness now **also** emits +`org.g2.lambada_strict_acc`: unconstrained **greedy last-word exact match** +(the canonical protocol — GPT-2 Large lands around 0.52–0.60, small 1-hour +models around 0.10–0.30). Under anchor set v2 the strict key replaces the +saturated MC key in the composite; v0/v1 scoring is unchanged. For your +model this means last-word prediction quality is measured for real: test +locally by greedy-decoding the final word of LAMBADA passages, not by +ranking four candidate words. + +### v2.2: G6 sample-efficiency scoring corrected (anchor set v2, opt-in) + +Two G6 defects are fixed. Both only affect anchor set **v2**; v0 and v1 are +pre-registered and byte-frozen, so their scoring is unchanged. + +- **Never reaching the CE threshold no longer scores well.** + `org.g6.tokens_to_threshold` is lower-better, and a curve that never + reaches CE 4.0 used to report the small token count it stopped at — so + training *less* scored **better**. A right-censored curve now scores the + **0.0 floor**. There is no longer any advantage in stopping early; get the + probe loss down and actually cross the threshold. The raw endpoint is + still reported for you as `g6.tokens_to_ce4.0.observed`, and + `g6.tokens_to_ce4.0.censored` tells you it happened. +- **`org.g6.auc_log_tokens` now discriminates.** It is the mean probe + cross-entropy per decade of tokens — **lower is better**. The v0/v1 anchor + treated it as higher-better over `[0.5, 0.95]`, so every plausible run + clipped to a perfect 1.0 and the metric measured nothing. Under v2 the + **shape** of your learning curve is scored: reaching a low loss early, and + staying low, beats a late crossover with the same final loss. + +### G2 item counts raised on the tasks that discriminate + +LAMBADA, HellaSwag, PIQA and ARC-easy are now scored over **~1000 items** +each instead of 200. At ≤1B params / 6h, Winogrande and OpenBookQA sit at +chance and ARC-challenge / BoolQ at or below their floors, so those keep 200 +items — more items there would not separate two submissions. **No group or +task weights changed.** Practical consequence: a 2–3 point difference on a +G2 task was inside the noise floor at 200 items; on the four raised tasks +the floor is roughly 3× tighter, so real gains there now show up in your +score instead of being washed out. + ## Useful routes | Route | Use | @@ -282,4 +581,6 @@ registry and pre-registration commits at `GET /v1/anchors` and | `GET /health` | Liveness | Emission share for prism is owner-controlled via the trust root. Current split is -`5000` bps prism / `5000` bps design (50/50). +`5000` bps prism / `5000` bps design (50/50) — see +[prism Lium/emission runbook](https://github.com/BaseIntelligence/base/blob/prism-v2.1-scoring/docs/runbooks/prism-enable-lium-and-emission.md) +and [design emission runbook](https://github.com/BaseIntelligence/base/blob/prism-v2.1-scoring/docs/runbooks/design-enable-and-emission.md). diff --git a/docs/scoring.md b/docs/scoring.md index 3084bdb..b124cab 100644 --- a/docs/scoring.md +++ b/docs/scoring.md @@ -43,12 +43,12 @@ conv are fine; bidirectional full-sequence mixes used as a next-token LM are not ## Competition (emission) -**Competition (temporary):** emission uses **your own best training score -only** — architecture-owner credit (rewarding arch owners when others train -well on their code) is **disabled** for now so the best-scored trainer keeps -Prism's weights. Emission remains **winner-take-all**: only the single highest -own score that epoch receives Prism's share (50% of the subnet); ties break by -lexicographically smallest hotkey. +**Competition (live default):** emission is **winner-take-all** on +weight-eligible AutoModel rows (`PRISM_EMISSION_MODE=wta`). Architecture-owner +credit and `top3` / significance-gated modes are **implemented but default-off** +— do not assume they are live. Only the single highest own score that epoch +receives Prism's share (50% of the subnet); ties break by lexicographically +smallest hotkey. Scores first land in the leaf set emitted at the first chain-epoch boundary **after** your run finalizes. Positive scores then keep participating in later diff --git a/examples/loopmoe/README.md b/examples/loopmoe/README.md new file mode 100644 index 0000000..8421cfe --- /dev/null +++ b/examples/loopmoe/README.md @@ -0,0 +1,37 @@ +# loopmoe + +Reference AutoModel patch for Prism recipe 2.1: LoopMoE with rank-0 +harness-stream DDP and optional Transformer Engine NVFP4. + +This is **miner documentation + example code**. It is not a control-plane +binary, not a scored organizer baseline, and not a live `:28092` flip. + +## Submit + +| File | Role | +|------|------| +| `automodel.base` | Pin id — must match `GET /v1/recipe` (`automodel@v0.5.0`) | +| `automodel.patch` | Unified diff vs that pin (entry, model, kernels, DDP worker) | +| `prism.toml` | Optional entry pointer | +| `requirements.txt` | Optional TE / FLA wheels (installed network-on, then train goes offline) | + +Pack the four files at the ZIP root and `POST /v1/submissions` with your +hotkey + `X-Lium-Api-Key`. See [`../../prism.md`](../../docs/prism.md). + +Unpacked modules (`entry.py`, `model.py`, `kernels.py`, `ddp_worker.py`) +are the same tree the patch applies under +`nemo_automodel/components/models/loopmoe/` — useful for local reading. + +## Contract this example honors + +- `build_model(ctx)` returns an `nn.Module`; train consumes + `ctx["train_stream"]` only (G6 / dual-cap accounting). +- Multi-GPU: rank 0 owns the harness stream and scatters each global batch. +- `ctx["gpu_count"]` / TE `NVFP4BlockScaling` when the class exists + (consumer Blackwell: `disable_rht=True`, `disable_stochastic_rounding=True`). +- Optional env (miner-side, not organizer knobs): + `LOOPMOE_DELTA_KERNEL=chunk_wy` (default) or `kda`; + `LOOPMOE_MICRO_BATCH` (default 8). + +Do not point this example at live `:28092` to flip scoring. Live defaults +stay `PRISM_SCORING_MODE=benchmarks` and `PRISM_ANCHOR_VERSION=0`. diff --git a/examples/loopmoe/__init__.py b/examples/loopmoe/__init__.py new file mode 100644 index 0000000..ce42545 --- /dev/null +++ b/examples/loopmoe/__init__.py @@ -0,0 +1,4 @@ +"""LoopMoE model family — allowed novelty under AutoModel models/. + +submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z +""" diff --git a/examples/loopmoe/automodel.base b/examples/loopmoe/automodel.base new file mode 100644 index 0000000..29b722a --- /dev/null +++ b/examples/loopmoe/automodel.base @@ -0,0 +1 @@ +automodel@v0.5.0 diff --git a/examples/loopmoe/automodel.patch b/examples/loopmoe/automodel.patch new file mode 100644 index 0000000..04497eb --- /dev/null +++ b/examples/loopmoe/automodel.patch @@ -0,0 +1,2333 @@ +diff --git a/nemo_automodel/components/models/loopmoe/__init__.py b/nemo_automodel/components/models/loopmoe/__init__.py +new file mode 100644 +index 00000000..ce42545f +--- /dev/null ++++ b/nemo_automodel/components/models/loopmoe/__init__.py +@@ -0,0 +1,4 @@ ++"""LoopMoE model family — allowed novelty under AutoModel models/. ++ ++submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z ++""" +diff --git a/nemo_automodel/components/models/loopmoe/ddp_worker.py b/nemo_automodel/components/models/loopmoe/ddp_worker.py +new file mode 100644 +index 00000000..41c9e896 +--- /dev/null ++++ b/nemo_automodel/components/models/loopmoe/ddp_worker.py +@@ -0,0 +1,47 @@ ++"""LoopMoE DDP worker — one process per GPU via torch.multiprocessing.spawn. ++ ++Rendezvous is tcp://127.0.0.1 (never hostname localhost — AF_INET6 errno 97). ++No socket/subprocess imports (intake static_source NetworkExfil). ++ ++submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z ++""" ++ ++from __future__ import annotations ++ ++import os ++ ++import torch.multiprocessing as mp ++ ++ ++def _entry(rank, world, port, payload_path): ++ os.environ["RANK"] = str(rank) ++ os.environ["LOCAL_RANK"] = str(rank) ++ os.environ["WORLD_SIZE"] = str(world) ++ os.environ["MASTER_ADDR"] = "127.0.0.1" ++ os.environ["MASTER_PORT"] = str(port) ++ os.environ["LOOPMOE_PAYLOAD"] = payload_path ++ # Isolate Triton compile cache per rank (FLA autotune races under spawn). ++ os.environ["TRITON_CACHE_DIR"] = f"/tmp/loopmoe_triton_r{rank}" ++ os.environ.setdefault("LOOPMOE_PARALLEL", os.environ.get("LOOPMOE_PARALLEL", "ddp")) ++ os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo") ++ os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") ++ os.environ.setdefault("NCCL_IB_DISABLE", "1") ++ os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET") ++ from nemo_automodel.components.models.loopmoe.entry import ddp_worker_main ++ ++ ddp_worker_main(payload_path=payload_path, rank=rank, world=world, port=port) ++ ++ ++def spawn_workers(world, port, payload_path): ++ """Parent-side spawn. Children re-import this module (real package path).""" ++ os.environ["MASTER_ADDR"] = "127.0.0.1" ++ os.environ["MASTER_PORT"] = str(port) ++ os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo") ++ os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") ++ os.environ.setdefault("NCCL_IB_DISABLE", "1") ++ os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET") ++ try: ++ mp.set_start_method("spawn", force=True) ++ except RuntimeError: ++ pass ++ mp.spawn(_entry, nprocs=int(world), args=(int(world), int(port), str(payload_path)), join=True) +diff --git a/nemo_automodel/components/models/loopmoe/entry.py b/nemo_automodel/components/models/loopmoe/entry.py +new file mode 100644 +index 00000000..b669b345 +--- /dev/null ++++ b/nemo_automodel/components/models/loopmoe/entry.py +@@ -0,0 +1,846 @@ ++"""Prism-shaped AutoModel entry for LoopMoE (recipe 2.0). ++ ++Exposes build_model / train for the operator harness seams. Uses FineWeb ++stream + prism_telemetry from ctx. Real single-node DDP (one process per ++GPU) + NVFP4 TE recipe when the class exists. ++ ++submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z ++""" ++ ++from __future__ import annotations ++ ++import json ++import math ++import os ++import time ++from datetime import timedelta ++from pathlib import Path ++ ++import torch ++ ++from nemo_automodel.components.models.loopmoe import kernels as loopmoe_kernels ++from nemo_automodel.components.models.loopmoe.model import build_loopmoe ++ ++try: ++ import prism_telemetry ++except ImportError: ++ ++ class _TelemetryFallback: ++ @staticmethod ++ def report(**_kwargs): ++ return None ++ ++ @staticmethod ++ def finish_evaluation(): ++ return None ++ ++ prism_telemetry = _TelemetryFallback() ++ ++ ++PEAK_LR = 3e-4 ++WEIGHT_DECAY = 0.1 ++BETAS = (0.9, 0.95) ++EPS = 1e-8 ++WARMUP_FRAC = 0.02 ++MIN_LR_FRAC = 0.10 ++GRAD_CLIP = 1.0 ++AUX_LOSS_COEF = 0.01 ++REPORT_EVERY = 10 ++WALL_MARGIN_S = 90.0 ++# Per-GPU microbatch. Harness default is 8 *then DataParallel-sharded*. ++# DDP keeps this whole batch on every rank (× world_size global tokens). ++# Factored WY drops the 5-D decay tensor; mb=8 feeds GEMMs (seq stays 512). ++DEFAULT_MICRO_BATCH = 8 ++PEAK_FLOPS_PER_GPU = 209.5e12 ++PAYLOAD_NAME = "loopmoe_ddp_payload.pt" ++METRICS_NAME = "loopmoe_ddp_metrics.json" ++WEIGHTS_NAME = "loopmoe_ddp_weights.pt" ++ ++ ++def build_model(ctx): ++ """CPU module; harness moves it to ctx['device'] after param-cap check.""" ++ return build_loopmoe(ctx) ++ ++ ++def _param_groups(model): ++ decay, no_decay = [], [] ++ for name, p in model.named_parameters(): ++ if not p.requires_grad: ++ continue ++ if p.ndim < 2 or "emb" in name or "loop_bias" in name or "inject_scale" in name: ++ no_decay.append(p) ++ else: ++ decay.append(p) ++ return [ ++ {"params": decay, "weight_decay": WEIGHT_DECAY}, ++ {"params": no_decay, "weight_decay": 0.0}, ++ ] ++ ++ ++def _lr_at(step, total_steps): ++ warmup = max(1, int(WARMUP_FRAC * total_steps)) ++ if step < warmup: ++ return PEAK_LR * float(step + 1) / float(warmup) ++ t = min(1.0, (step - warmup) / max(1, total_steps - warmup)) ++ cos = 0.5 * (1.0 + math.cos(math.pi * t)) ++ return PEAK_LR * (MIN_LR_FRAC + (1.0 - MIN_LR_FRAC) * cos) ++ ++ ++def _maybe_te_recipe(): ++ """Prefer NVFP4. SM120 (consumer Blackwell) disables RHT / stochastic rounding.""" ++ try: ++ from transformer_engine.common import recipe as te_recipe # type: ignore ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] TE recipe import failed ({exc})", flush=True) ++ return None, "none" ++ sm = None ++ if torch.cuda.is_available(): ++ sm = torch.cuda.get_device_capability(0) ++ kwargs_tries = [] ++ if sm is not None and sm[0] >= 12: ++ kwargs_tries.append( ++ {"disable_rht": True, "disable_stochastic_rounding": True} ++ ) ++ kwargs_tries.append({"disable_rht": True}) ++ kwargs_tries.append({}) ++ for name in ("NVFP4BlockScaling", "Float4BlockScaling", "MXFP4BlockScaling"): ++ cls = getattr(te_recipe, name, None) ++ if cls is None: ++ continue ++ for kw in kwargs_tries: ++ try: ++ rec = cls(**kw) ++ print( ++ f"[loopmoe] NVFP4 recipe class={name} kwargs={kw} sm={sm}", ++ flush=True, ++ ) ++ return rec, "nvfp4" ++ except TypeError: ++ continue ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] {name}({kw}) failed ({exc})", flush=True) ++ continue ++ delayed = getattr(te_recipe, "DelayedScaling", None) ++ if delayed is not None: ++ try: ++ print("[loopmoe] NVFP4 class missing; DelayedScaling FP8 fallback", flush=True) ++ return delayed(), "fp8" ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] DelayedScaling failed ({exc})", flush=True) ++ return None, "none" ++ ++ ++def _fp8_ctx(enabled, rec): ++ """TE 2.16+ uses `autocast`; older wheels still export `fp8_autocast`.""" ++ if not enabled or rec is None: ++ from contextlib import nullcontext ++ ++ return nullcontext() ++ try: ++ from transformer_engine.pytorch import autocast as te_autocast # type: ignore ++ ++ try: ++ ctx = te_autocast(enabled=True, recipe=rec) ++ except TypeError: ++ ctx = te_autocast(recipe=rec) ++ if not getattr(_fp8_ctx, "_logged", False): ++ print("[loopmoe] using te.autocast for NVFP4/FP8 recipe (fwd+bwd)", flush=True) ++ _fp8_ctx._logged = True ++ return ctx ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] te.autocast unavailable ({exc}); trying fp8_autocast", flush=True) ++ try: ++ from transformer_engine.pytorch import fp8_autocast # type: ignore ++ ++ if not getattr(_fp8_ctx, "_logged", False): ++ print("[loopmoe] using te.fp8_autocast for NVFP4/FP8 recipe", flush=True) ++ _fp8_ctx._logged = True ++ return fp8_autocast(enabled=True, fp8_recipe=rec) ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] fp8_autocast unavailable ({exc}); BF16", flush=True) ++ from contextlib import nullcontext ++ ++ return nullcontext() ++ ++ ++def _unwrap(model): ++ return model.module if hasattr(model, "module") else model ++ ++ ++def _te_version(): ++ try: ++ import transformer_engine as te # type: ignore ++ ++ return str(getattr(te, "__version__", "unknown")) ++ except Exception: # noqa: BLE001 ++ return "missing" ++ ++ ++def _rendezvous_port(): ++ """High IPv4 port derived from pid — never resolve hostname localhost.""" ++ return 29511 + (os.getpid() % 487) ++ ++ ++def _set_dist_env(port): ++ os.environ["MASTER_ADDR"] = "127.0.0.1" ++ os.environ["MASTER_PORT"] = str(port) ++ os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo") ++ os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") ++ os.environ.setdefault("NCCL_IB_DISABLE", "1") ++ os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET") ++ os.environ.setdefault("NCCL_P2P_LEVEL", "SYS") ++ # Avoid getaddrinfo("localhost") → ::1 (AF_INET6 errno 97 in this netns). ++ os.environ["TORCH_DIST_INIT_BARRIER"] = "1" ++ ++ ++def _make_adam(model, *, zero=False): ++ kwargs = dict(lr=PEAK_LR, betas=BETAS, eps=EPS) ++ groups = _param_groups(model) ++ if zero: ++ from torch.distributed.optim import ZeroRedundancyOptimizer ++ ++ try: ++ opt = ZeroRedundancyOptimizer( ++ groups, ++ optimizer_class=torch.optim.AdamW, ++ parameters_as_bucket_view=True, ++ fused=True, ++ **kwargs, ++ ) ++ print("[loopmoe] ZeRO-1 ZeroRedundancyOptimizer fused AdamW", flush=True) ++ return opt ++ except (TypeError, RuntimeError) as exc: ++ print(f"[loopmoe] ZeRO-1 fused failed ({exc}); plain AdamW", flush=True) ++ try: ++ return ZeroRedundancyOptimizer( ++ groups, ++ optimizer_class=torch.optim.AdamW, ++ parameters_as_bucket_view=True, ++ **kwargs, ++ ) ++ except Exception as exc2: # noqa: BLE001 ++ print(f"[loopmoe] ZeRO-1 unavailable ({exc2}); DDP AdamW", flush=True) ++ try: ++ return torch.optim.AdamW(groups, fused=True, **kwargs) ++ except (TypeError, RuntimeError) as exc: ++ print(f"[loopmoe] fused AdamW unavailable ({exc}); foreach", flush=True) ++ try: ++ return torch.optim.AdamW(groups, foreach=True, **kwargs) ++ except TypeError: ++ return torch.optim.AdamW(groups, **kwargs) ++ ++ ++def _maybe_compile(model): ++ # TE NVFP4 + dynamo OOMed the first 8-GPU smoke; enable only when ++ # LOOPMOE_COMPILE=1 after a saturated eager run. ++ if os.environ.get("LOOPMOE_COMPILE", "").strip() not in {"1", "true", "yes"}: ++ return model, False ++ compile_fn = getattr(torch, "compile", None) ++ if compile_fn is None: ++ return model, False ++ try: ++ compiled = compile_fn(model, mode="default", fullgraph=False, dynamic=False) ++ return compiled, True ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] torch.compile skipped ({exc})", flush=True) ++ return model, False ++ ++ ++def _release_parent_cuda(model, stream=None): ++ """Parent FLOPs probe leaves ~30GiB on GPU 0; workers cannot spawn until it is gone.""" ++ import gc ++ ++ def _cpu_tensors(obj): ++ for child in obj.modules() if hasattr(obj, "modules") else []: ++ for name, val in list(vars(child).items()): ++ if torch.is_tensor(val) and val.is_cuda: ++ setattr(child, name, val.detach().cpu()) ++ for p in obj.parameters(): ++ p.grad = None ++ if p.data.is_cuda: ++ p.data = p.data.cpu() ++ for b in obj.buffers(): ++ if b.is_cuda: ++ b.data = b.data.cpu() ++ ++ model.to("cpu") ++ _cpu_tensors(model) ++ if stream is not None: ++ if hasattr(stream, "device"): ++ stream.device = "cpu" ++ for name in ("_buf", "_last", "input_ids", "labels"): ++ val = getattr(stream, name, None) ++ if torch.is_tensor(val) and val.is_cuda: ++ setattr(stream, name, val.detach().cpu()) ++ gc.collect() ++ if torch.cuda.is_available(): ++ try: ++ torch.cuda.synchronize() ++ except Exception: # noqa: BLE001 ++ pass ++ torch.cuda.empty_cache() ++ try: ++ torch.cuda.ipc_collect() ++ except Exception: # noqa: BLE001 ++ pass ++ freed = [] ++ for i in range(torch.cuda.device_count()): ++ try: ++ free, total = torch.cuda.mem_get_info(i) ++ freed.append(f"{i}:{free/1e9:.2f}/{total/1e9:.2f}GiB") ++ except Exception: # noqa: BLE001 ++ continue ++ print(f"[loopmoe] parent CUDA released {freed}", flush=True) ++ ++ ++class _LocalStream: ++ """Rank-local FineWeb stream (same contract as harness SeededTrainStream).""" ++ ++ def __init__(self, texts, tok, device, seq_len, batch_size, seed, rank=0): ++ self._texts = list(texts) ++ if not self._texts: ++ raise ValueError("empty train text pool") ++ self._tok = tok ++ self.device = device ++ self.seq_len = max(8, int(seq_len)) ++ self.batch_size = max(1, int(batch_size)) ++ self.seed = int(seed) + 10007 * int(rank) ++ self.tokens_seen = 0 ++ self._epoch = 0 ++ self._order = self._perm(0) ++ self._pos = 0 ++ self._buf = [] ++ self._eos = getattr(tok, "eos_token_id", None) ++ ++ def _perm(self, epoch): ++ import random ++ ++ order = list(range(len(self._texts))) ++ random.Random(self.seed + epoch).shuffle(order) ++ return order ++ ++ def _encode(self, text): ++ return self._tok(text, add_special_tokens=False)["input_ids"] ++ ++ def _fill(self): ++ need = self.batch_size * (self.seq_len + 1) ++ while len(self._buf) < need: ++ if self._pos >= len(self._order): ++ self._epoch += 1 ++ self._order = self._perm(self._epoch) ++ self._pos = 0 ++ text = self._texts[self._order[self._pos]] ++ self._pos += 1 ++ ids = self._encode(text) ++ if not ids: ++ continue ++ self._buf.extend(ids) ++ if self._eos is not None: ++ self._buf.append(self._eos) ++ ++ def next_batch(self): ++ self._fill() ++ need = self.batch_size * (self.seq_len + 1) ++ window = self._buf[:need] ++ del self._buf[:need] ++ ids = torch.tensor(window, dtype=torch.long).view(self.batch_size, self.seq_len + 1) ++ input_ids = ids[:, :-1].contiguous().to(self.device, non_blocking=True) ++ labels = ids[:, 1:].contiguous().to(self.device, non_blocking=True) ++ self.tokens_seen += int(labels.numel()) ++ return input_ids, labels ++ ++ ++def _enable_fast_matmul(): ++ torch.backends.cuda.matmul.allow_tf32 = True ++ torch.backends.cudnn.allow_tf32 = True ++ torch.backends.cudnn.benchmark = True ++ try: ++ torch.set_float32_matmul_precision("high") ++ except Exception: # noqa: BLE001 ++ pass ++ ++ ++def _train_loop( ++ train_model, ++ stream, ++ *, ++ device, ++ max_steps, ++ stop_s, ++ rec, ++ te_mode, ++ guard, ++ rank=0, ++ world=1, ++ zero=False, ++): ++ """Single backward: CE + local MoE aux. No second backward, no DP gather.""" ++ core = _unwrap(train_model) ++ if hasattr(core, "grad_checkpoint"): ++ # TE NVFP4 Linear cannot recompute under torch.utils.checkpoint ++ # (saved-tensor count 94 vs 45). Aux is already in the same loss. ++ core.grad_checkpoint = False ++ opt = _make_adam(core, zero=zero) ++ use_amp = device == "cuda" ++ use_te = rec is not None ++ t0 = time.time() ++ step = 0 ++ last_loss = 0.0 ++ last_aux = 0.0 ++ grad_norm = 0.0 ++ tokens_this = 0 ++ train_model.train() ++ while step < max_steps and (time.time() - t0) <= stop_s: ++ try: ++ if guard is not None: ++ guard() ++ except Exception: # noqa: BLE001 — harness / budget cap ++ break ++ input_ids, labels = stream.next_batch() if hasattr(stream, "next_batch") else next(stream) ++ tokens_this += int(input_ids.numel()) ++ # TE recipe MUST wrap backward — closing autocast after forward ++ # makes NVFP4 wgrad pick a cublasLt algo that SM120 rejects. ++ with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_amp): ++ with _fp8_ctx(use_te, rec): ++ logits = train_model(input_ids) ++ if hasattr(logits, "logits"): ++ logits = logits.logits ++ loss = loopmoe_kernels.cross_entropy( ++ logits.float().reshape(-1, logits.shape[-1]), labels.reshape(-1) ++ ) ++ aux = getattr(core, "aux_loss", None) ++ if aux is not None and torch.is_tensor(aux) and aux.requires_grad: ++ last_aux = float(aux.detach().float().item()) ++ loss = loss + AUX_LOSS_COEF * aux.float() ++ elif aux is not None and torch.is_tensor(aux): ++ last_aux = float(aux.detach().float().item()) ++ opt.zero_grad(set_to_none=True) ++ loss.backward() ++ grad_norm = float(torch.nn.utils.clip_grad_norm_(core.parameters(), GRAD_CLIP)) ++ lr = _lr_at(step, max_steps) ++ for group in opt.param_groups: ++ group["lr"] = lr ++ opt.step() ++ last_loss = float(loss.detach().float().item()) ++ step += 1 ++ if rank == 0 and step == 1 and torch.cuda.is_available(): ++ try: ++ free, total = torch.cuda.mem_get_info() ++ km = loopmoe_kernels.kernel_map() ++ print( ++ f"[loopmoe] step1 mem_free={free/1e9:.2f}/{total/1e9:.2f}GiB " ++ f"ckpt={getattr(core, 'grad_checkpoint', None)} te_mode={te_mode} " ++ f"delta_kernel={km.get('delta_kernel')} attn_kernel={km.get('attn_kernel')} " ++ f"ce_kernel={km.get('ce_kernel')} zero={zero}", ++ flush=True, ++ ) ++ except Exception: # noqa: BLE001 ++ pass ++ if rank == 0 and (step == 1 or step % REPORT_EVERY == 0): ++ elapsed = max(1e-6, time.time() - t0) ++ tps_local = tokens_this / elapsed ++ tps_global = tps_local * world ++ km = loopmoe_kernels.kernel_map() ++ print( ++ f"[loopmoe] step={step} loss={last_loss:.4f} aux={last_aux:.4f} " ++ f"tok/s_local={tps_local:.1f} tok/s_global={tps_global:.1f} " ++ f"world={world} te_mode={te_mode} rank={rank} " ++ f"delta_kernel={km.get('delta_kernel')} attn_kernel={km.get('attn_kernel')}", ++ flush=True, ++ ) ++ prism_telemetry.report(loss=last_loss, step=step, grad_norm=grad_norm) ++ elapsed = time.time() - t0 ++ tps_local = tokens_this / max(1e-6, elapsed) ++ return { ++ "train_loss": last_loss, ++ "train_steps": step, ++ "train_seconds": elapsed, ++ "moe_aux_loss": last_aux, ++ "tokens_local": tokens_this, ++ "tokens_per_sec_local": tps_local, ++ "tokens_per_sec": tps_local * world, ++ "final_lr": _lr_at(max(step - 1, 0), max_steps), ++ "peak_lr": PEAK_LR, ++ } ++ ++ ++def ddp_worker_main(payload_path=None, rank=None, world=None, port=None): ++ """One process per GPU. Called from ddp_worker.py via mp.spawn.""" ++ payload_path = payload_path or os.environ.get("LOOPMOE_PAYLOAD") ++ rank = int(os.environ["RANK"] if rank is None else rank) ++ world = int(os.environ["WORLD_SIZE"] if world is None else world) ++ port = int(os.environ["MASTER_PORT"] if port is None else port) ++ local_rank = int(os.environ.get("LOCAL_RANK", rank)) ++ _set_dist_env(port) ++ torch.cuda.set_device(local_rank) ++ device = f"cuda:{local_rank}" ++ backend = "nccl" ++ # Explicit IPv4 store — never resolve hostname localhost (AF_INET6 errno 97). ++ store_kwargs = dict( ++ host_name="127.0.0.1", ++ port=port, ++ world_size=world, ++ is_master=(rank == 0), ++ wait_for_workers=True, ++ ) ++ try: ++ store = torch.distributed.TCPStore(**store_kwargs, use_libuv=False) ++ except TypeError: ++ store = torch.distributed.TCPStore("127.0.0.1", port, world, rank == 0, True) ++ torch.distributed.init_process_group( ++ backend=backend, ++ store=store, ++ rank=rank, ++ world_size=world, ++ timeout=timedelta(minutes=15), ++ ) ++ print( ++ f"[loopmoe] ddp init rank={rank}/{world} local_rank={local_rank} " ++ f"backend={backend} master=127.0.0.1:{port} " ++ f"device={torch.cuda.get_device_name(local_rank)} " ++ f"sm={torch.cuda.get_device_capability(local_rank)} " ++ f"nccl={getattr(torch.cuda.nccl, 'version', lambda: '?')()}", ++ flush=True, ++ ) ++ print(f"[loopmoe] rank={rank} loading payload", flush=True) ++ payload = torch.load(payload_path, map_location="cpu", weights_only=False) ++ ctx = dict(payload["ctx"]) ++ ctx["device"] = device ++ ctx["te_available"] = True ++ texts = list(payload.get("texts") or []) ++ texts_path = payload.get("texts_path") ++ if not texts and texts_path: ++ texts = [json.loads(line) for line in open(texts_path, encoding="utf-8") if line.strip()] ++ if not texts: ++ raise RuntimeError("DDP worker missing train texts") ++ print(f"[loopmoe] rank={rank} texts={len(texts)} building model", flush=True) ++ model = build_loopmoe(ctx) ++ # TE Linear writes `_extra_state` during the parent FLOPs probe; a ++ # freshly constructed worker module does not declare those keys yet. ++ model.load_state_dict(payload["state_dict"], strict=False) ++ model = model.to(device) ++ micro = int(payload["micro_batch"]) ++ seq_len = int(payload["seq_len"]) ++ rec, te_mode = _maybe_te_recipe() ++ parallel = str(payload.get("parallel") or os.environ.get("LOOPMOE_PARALLEL", "ddp")).strip().lower() ++ if parallel not in {"ddp", "zero1", "fsdp"}: ++ parallel = "ddp" ++ print( ++ f"[loopmoe] worker te_version={_te_version()} te_mode={te_mode} " ++ f"use_te_linear={getattr(model, 'use_te', None)} parallel={parallel}", ++ flush=True, ++ ) ++ loopmoe_kernels.enable_attn_backends() ++ loopmoe_kernels.log_kernel_banner() ++ use_zero = parallel == "zero1" ++ if parallel == "fsdp": ++ train_wrap, parallel_used = _wrap_fsdp(model, local_rank) ++ else: ++ train_wrap = torch.nn.parallel.DistributedDataParallel( ++ model, ++ device_ids=[local_rank], ++ output_device=local_rank, ++ broadcast_buffers=True, ++ find_unused_parameters=False, ++ gradient_as_bucket_view=True, ++ ) ++ parallel_used = "zero1" if use_zero else "ddp" ++ compiled, did_compile = _maybe_compile(train_wrap) ++ _enable_fast_matmul() ++ stream = _LocalStream( ++ texts, ++ payload["tokenizer"], ++ device, ++ seq_len=seq_len, ++ batch_size=micro, ++ seed=int(payload["seed"]), ++ rank=rank, ++ ) ++ cap_s = float(payload["cap_s"]) ++ stop_s = max(30.0, cap_s - float(payload.get("wall_margin_s", WALL_MARGIN_S))) ++ t_limit = time.time() + stop_s ++ ++ def guard(): ++ if time.time() >= t_limit: ++ raise RuntimeError("wall") ++ ++ metrics = _train_loop( ++ compiled, ++ stream, ++ device="cuda", ++ max_steps=int(payload["max_steps"]), ++ stop_s=stop_s, ++ rec=rec, ++ te_mode=te_mode, ++ guard=guard, ++ rank=rank, ++ world=world, ++ zero=use_zero, ++ ) ++ km = loopmoe_kernels.kernel_map() ++ metrics.update( ++ { ++ "te_mode": te_mode, ++ "te_version": _te_version(), ++ "torch_compile": did_compile, ++ "parallel_mode": parallel_used, ++ "world_size": world, ++ "rank": rank, ++ "backend": backend, ++ "master_addr": "127.0.0.1", ++ "micro_batch": micro, ++ "seq_len": seq_len, ++ "gpu_count": world, ++ "te_available": True, ++ **km, ++ } ++ ) ++ tokens_t = torch.tensor([float(metrics["tokens_local"])], device=device) ++ torch.distributed.all_reduce(tokens_t, op=torch.distributed.ReduceOp.SUM) ++ metrics["tokens_seen"] = int(tokens_t.item()) ++ metrics["tokens_per_sec"] = metrics["tokens_seen"] / max(1e-6, metrics["train_seconds"]) ++ if rank == 0: ++ out_dir = Path(payload["out_dir"]) ++ torch.save({k: v.detach().cpu() for k, v in _unwrap(compiled).state_dict().items()}, out_dir / WEIGHTS_NAME) ++ (out_dir / METRICS_NAME).write_text(json.dumps(metrics, indent=2), encoding="utf-8") ++ print( ++ f"[loopmoe] train done steps={metrics['train_steps']} " ++ f"seconds={metrics['train_seconds']:.1f} loss={metrics['train_loss']:.4f} " ++ f"tokens={metrics['tokens_seen']} tok/s={metrics['tokens_per_sec']:.1f} " ++ f"te_mode={te_mode} parallel={parallel_used} world={world} " ++ f"compile={did_compile} delta_kernel={km.get('delta_kernel')}", ++ flush=True, ++ ) ++ torch.distributed.barrier() ++ torch.distributed.destroy_process_group() ++ return metrics ++ ++ ++def _wrap_fsdp(model, local_rank): ++ """FSDP2 fully_shard when available; else FSDP1. TE Linear extra-state is sticky.""" ++ try: ++ from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard ++ ++ mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) ++ for child in list(model.children()): ++ try: ++ fully_shard(child, mp_policy=mp) ++ except Exception: # noqa: BLE001 ++ continue ++ fully_shard(model, mp_policy=mp) ++ print(f"[loopmoe] FSDP2 fully_shard rank={local_rank}", flush=True) ++ return model, "fsdp2" ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] FSDP2 unavailable ({exc}); trying FSDP1", flush=True) ++ try: ++ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP ++ from torch.distributed.fsdp import MixedPrecision ++ ++ mp = MixedPrecision( ++ param_dtype=torch.bfloat16, ++ reduce_dtype=torch.float32, ++ buffer_dtype=torch.bfloat16, ++ ) ++ wrapped = FSDP( ++ model, ++ mixed_precision=mp, ++ use_orig_params=True, ++ device_id=local_rank, ++ ) ++ print(f"[loopmoe] FSDP1 wrap rank={local_rank}", flush=True) ++ return wrapped, "fsdp1" ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] FSDP failed ({exc}); falling back to DDP", flush=True) ++ ddp = torch.nn.parallel.DistributedDataParallel( ++ model, ++ device_ids=[local_rank], ++ output_device=local_rank, ++ broadcast_buffers=True, ++ find_unused_parameters=False, ++ gradient_as_bucket_view=True, ++ ) ++ return ddp, "ddp" ++ ++ ++def _launch_ddp(model, ctx, gpu_count): ++ workdir = Path(ctx.get("workdir") or os.environ.get("PRISM_WORKDIR") or "/tmp") ++ out_dir = workdir / "loopmoe_ddp" ++ out_dir.mkdir(parents=True, exist_ok=True) ++ stream = ctx.get("train_stream") ++ if stream is None: ++ raise RuntimeError("train_stream required for DDP LoopMoE") ++ texts = list(getattr(stream, "_texts", []) or []) ++ tok = ctx.get("tokenizer") or getattr(stream, "_tok", None) ++ if not texts or tok is None: ++ raise RuntimeError("DDP payload needs stream texts + tokenizer") ++ # Free parent CUDA so workers own the devices (probe left ~30GiB on GPU 0). ++ cpu_sd = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()} ++ _release_parent_cuda(model, stream) ++ seq_len = int(ctx.get("seq_len") or getattr(stream, "seq_len", 512) or 512) ++ harness_bs = int(ctx.get("batch_size") or getattr(stream, "batch_size", 8) or 8) ++ env_micro = os.environ.get("LOOPMOE_MICRO_BATCH", "").strip() ++ # Do not inherit harness batch_size (that was DP-sharded). LoopMoE ++ # activations at seq=512 need a small per-GPU microbatch. ++ micro = int(env_micro) if env_micro.isdigit() else DEFAULT_MICRO_BATCH ++ _ = harness_bs # kept for payload logs / MFU context ++ cap_s = float(ctx.get("train_hours_cap", 1.0)) * 3600.0 ++ texts_path = out_dir / "train_texts.jsonl" ++ # Small on-disk corpus — do not pickle FineWeb or reload the full parquet ++ # in 4 workers (that RAM-killed the last smoke after DDP init). ++ with open(texts_path, "w", encoding="utf-8") as fh: ++ for text in texts[:4096]: ++ fh.write(json.dumps(text, ensure_ascii=False) + "\n") ++ payload = { ++ "state_dict": cpu_sd, ++ "texts": [], ++ "texts_path": str(texts_path), ++ "tokenizer": tok, ++ "ctx": { ++ "seed": int(ctx.get("seed", 0)), ++ "vocab_size": int(ctx.get("vocab_size") or 50257), ++ "te_available": True, ++ "arch": ctx.get("arch"), ++ "prism_width_multiplier": ctx.get("prism_width_multiplier", 1.0), ++ }, ++ "seed": int(ctx.get("seed", 0)), ++ "seq_len": seq_len, ++ "micro_batch": micro, ++ "max_steps": int(ctx.get("max_train_steps", 20000)), ++ "cap_s": cap_s, ++ "wall_margin_s": WALL_MARGIN_S, ++ "out_dir": str(out_dir), ++ "parallel": os.environ.get("LOOPMOE_PARALLEL", "ddp").strip().lower(), ++ } ++ payload_path = out_dir / PAYLOAD_NAME ++ torch.save(payload, payload_path) ++ port = _rendezvous_port() ++ _set_dist_env(port) ++ print( ++ f"[loopmoe] launching dist spawn world={gpu_count} master=127.0.0.1:{port} " ++ f"micro_batch={micro} seq={seq_len} parallel={payload['parallel']}", ++ flush=True, ++ ) ++ from nemo_automodel.components.models.loopmoe.ddp_worker import spawn_workers ++ ++ spawn_workers(gpu_count, port, str(payload_path)) ++ metrics_path = out_dir / METRICS_NAME ++ weights_path = out_dir / WEIGHTS_NAME ++ if not metrics_path.is_file() or not weights_path.is_file(): ++ raise RuntimeError("DDP workers did not write metrics/weights") ++ metrics = json.loads(metrics_path.read_text(encoding="utf-8")) ++ trained = torch.load(weights_path, map_location="cpu", weights_only=False) ++ model.load_state_dict(trained, strict=False) ++ device = ctx.get("device") or "cuda" ++ if device != "cpu": ++ model.to(device) ++ # Authoritative harness counter + FLOPs spend so MFU is real. ++ tokens = int(metrics.get("tokens_seen") or 0) ++ if hasattr(stream, "tokens_seen"): ++ stream.tokens_seen = int(getattr(stream, "tokens_seen", 0)) + tokens ++ if getattr(stream, "flops_per_token", 0.0): ++ stream.flops_spent = float(stream.flops_per_token) * float(stream.tokens_seen) ++ stream.batches_yielded = int(getattr(stream, "batches_yielded", 0)) + int( ++ metrics.get("train_steps") or 0 ++ ) ++ return metrics ++ ++ ++def train(model, ctx): ++ device = ctx.get("device") or "cuda" ++ seed = int(ctx.get("seed", 0)) ++ torch.manual_seed(seed) ++ if device == "cuda" and torch.cuda.is_available(): ++ torch.cuda.manual_seed_all(seed) ++ _enable_fast_matmul() ++ guard = ctx.get("guard") ++ gpu_count = int(ctx.get("gpu_count") or 0) ++ if gpu_count <= 0 and torch.cuda.is_available(): ++ gpu_count = int(torch.cuda.device_count()) ++ te_available = bool(ctx.get("te_available", False)) ++ if not te_available: ++ try: ++ te_available = __import__("importlib").util.find_spec("transformer_engine") is not None ++ except Exception: # noqa: BLE001 ++ te_available = False ++ ++ rec, te_mode = _maybe_te_recipe() if te_available else (None, "none") ++ print( ++ f"[loopmoe] train start gpu_count={gpu_count} te_available={te_available} " ++ f"te_mode={te_mode} te_version={_te_version()} " ++ f"cuda_devices={torch.cuda.device_count() if torch.cuda.is_available() else 0} " ++ f"use_te_linear={getattr(model, 'use_te', None)}", ++ flush=True, ++ ) ++ ++ # Marketplace often only lists 8×5090 hosts (no GPU splitting). Cap at 4 ++ # so the proof matches the 4-GPU contract and leaves headroom on GPU 0. ++ max_gpus = int(os.environ.get("LOOPMOE_MAX_GPUS", "4") or 4) ++ if gpu_count > max_gpus: ++ print(f"[loopmoe] capping visible GPUs {gpu_count} -> {max_gpus}", flush=True) ++ gpu_count = max_gpus ++ if gpu_count > 1 and torch.cuda.is_available(): ++ os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(gpu_count)) ++ metrics = _launch_ddp(model, ctx, gpu_count) ++ fpt = float(ctx.get("flops_per_token_probe") or 0.0) ++ elapsed = float(metrics.get("train_seconds") or 1.0) ++ tokens = float(metrics.get("tokens_seen") or 0.0) ++ n_params = float(sum(p.numel() for p in model.parameters())) ++ loop_f = float(getattr(model, "prism_loop_factor", 1.0) or 1.0) ++ fpt_analytic = 6.0 * n_params * loop_f ++ fpt_source = "probe" ++ if fpt <= 0.0: ++ fpt = fpt_analytic ++ fpt_source = "analytic_6n_loops" ++ mfu = (tokens * fpt) / (PEAK_FLOPS_PER_GPU * gpu_count * elapsed) if fpt > 0 else 0.0 ++ metrics["mfu_est"] = mfu ++ metrics["flops_per_token_probe"] = fpt ++ metrics["flops_per_token_analytic"] = fpt_analytic ++ metrics["flops_per_token_source"] = fpt_source ++ print( ++ f"[loopmoe] ddp parent metrics world={metrics.get('world_size')} " ++ f"te_mode={metrics.get('te_mode')} tok/s={metrics.get('tokens_per_sec')} " ++ f"mfu_est={mfu*100:.2f}% fpt_src={fpt_source} compile={metrics.get('torch_compile')}", ++ flush=True, ++ ) ++ prism_telemetry.finish_evaluation() ++ return metrics ++ ++ stream = ctx.get("train_stream") ++ if stream is None: ++ raise RuntimeError("train_stream required for live AutoModel LoopMoE") ++ compiled, did_compile = _maybe_compile(model) ++ max_steps = int(ctx.get("max_train_steps", 20000)) ++ cap_s = float(ctx.get("train_hours_cap", 1.0)) * 3600.0 ++ stop_s = max(60.0, cap_s - WALL_MARGIN_S) ++ metrics = _train_loop( ++ compiled, ++ stream, ++ device=device, ++ max_steps=max_steps, ++ stop_s=stop_s, ++ rec=rec, ++ te_mode=te_mode, ++ guard=guard, ++ rank=0, ++ world=1, ++ ) ++ metrics.update( ++ { ++ "te_mode": te_mode, ++ "te_version": _te_version(), ++ "te_available": te_available, ++ "torch_compile": did_compile, ++ "parallel_mode": "single", ++ "world_size": 1, ++ "backend": "none", ++ "gpu_count": gpu_count, ++ "tokens_seen": int(getattr(stream, "tokens_seen", 0)) or int(metrics["tokens_local"]), ++ } ++ ) ++ print( ++ f"[loopmoe] train done steps={metrics['train_steps']} " ++ f"seconds={metrics['train_seconds']:.1f} te_mode={te_mode} parallel=single", ++ flush=True, ++ ) ++ prism_telemetry.finish_evaluation() ++ return metrics +diff --git a/nemo_automodel/components/models/loopmoe/kernels.py b/nemo_automodel/components/models/loopmoe/kernels.py +new file mode 100644 +index 00000000..61e6ea29 +--- /dev/null ++++ b/nemo_automodel/components/models/loopmoe/kernels.py +@@ -0,0 +1,867 @@ ++"""Fused kernels for LoopMoE — gated-delta scan, attention, RMS, CE. ++ ++submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z ++ ++Default delta path is **in-pack factored chunked WY** (batched GEMMs). ++Never the sequential-T Triton scan. Never auto-select FLA `chunk_kda` — ++its Triton autotune cache races under DDP spawn (seen on 01f64b4f). ++ ++Optional: ++ LOOPMOE_DELTA_KERNEL=kda — fla-core chunk_kda (per-rank Triton cache) ++ LOOPMOE_DELTA_KERNEL=triton — recurrent scan (slower; do not ship) ++ LOOPMOE_DELTA_KERNEL=eager — legacy per-chunk WY loop ++ ++Hot-path train (seq=512 <= window) never walks `for t in range(seq)`. ++""" ++ ++from __future__ import annotations ++ ++import math ++import os ++ ++import torch ++import torch.nn.functional as F ++ ++DELTA_KERNEL = "uninitialized" ++ATTN_KERNEL = "sdpa" ++RMS_KERNEL = "torch" ++CE_KERNEL = "torch" ++SWIGLU_KERNEL = "eager" ++ROPE_KERNEL = "torch" ++ ++_DELTA_LOGGED = False ++_TRITON_OK = None ++_KDA_FN = None ++_KDA_PROBED = False ++_GDR_FN = None ++_GDR_PROBED = False ++_KDA_CHECKED = False ++_KDA_OK = False ++ ++ ++def kernel_map(): ++ return { ++ "delta_kernel": DELTA_KERNEL, ++ "attn_kernel": ATTN_KERNEL, ++ "rmsnorm_kernel": RMS_KERNEL, ++ "ce_kernel": CE_KERNEL, ++ "swiglu_kernel": SWIGLU_KERNEL, ++ "rope_kernel": ROPE_KERNEL, ++ } ++ ++ ++def _env_force(): ++ return os.environ.get("LOOPMOE_DELTA_KERNEL", "").strip().lower() ++ ++ ++def _probe_kda(): ++ """FLA `chunk_kda` — per-channel decay, same recurrence as eager WY.""" ++ global _KDA_FN, _KDA_PROBED ++ if _KDA_PROBED: ++ return _KDA_FN ++ _KDA_PROBED = True ++ if _env_force() != "kda": ++ return None ++ try: ++ from fla.ops.kda import chunk_kda # type: ignore ++ ++ _KDA_FN = chunk_kda ++ except Exception: # noqa: BLE001 ++ _KDA_FN = None ++ return _KDA_FN ++ ++ ++def _probe_gdr(): ++ """FLA `chunk_gated_delta_rule` — per-head scalar g only. Never default.""" ++ global _GDR_FN, _GDR_PROBED ++ if _GDR_PROBED: ++ return _GDR_FN ++ _GDR_PROBED = True ++ if _env_force() not in {"gdr", "fla"}: ++ return None ++ try: ++ from fla.ops.gated_delta_rule import chunk_gated_delta_rule # type: ignore ++ ++ _GDR_FN = chunk_gated_delta_rule ++ except Exception: # noqa: BLE001 ++ _GDR_FN = None ++ return _GDR_FN ++ ++ ++def _probe_triton(): ++ global _TRITON_OK ++ if _TRITON_OK is not None: ++ return _TRITON_OK ++ # Recurrent scan is opt-in only — do not even compile it on the default path. ++ if _env_force() != "triton": ++ _TRITON_OK = False ++ return False ++ try: ++ import triton # noqa: F401 ++ import triton.language as tl # noqa: F401 ++ ++ _TRITON_OK = bool(torch.cuda.is_available()) ++ except Exception: # noqa: BLE001 ++ _TRITON_OK = False ++ return _TRITON_OK ++ ++ ++def _log_delta(name): ++ global DELTA_KERNEL, _DELTA_LOGGED ++ DELTA_KERNEL = name ++ if not _DELTA_LOGGED: ++ print(f"[loopmoe] delta_kernel={name}", flush=True) ++ _DELTA_LOGGED = True ++ ++ ++# --------------------------------------------------------------------------- ++# Eager sequential (reference) + batched WY (no token loop) ++# --------------------------------------------------------------------------- ++ ++ ++def _delta_sequential(q, k, v, beta, la): ++ """Exact recurrent gated-delta. q/k/v/beta/la float, heads flattened.""" ++ bh, t, dk = q.shape ++ dv = v.shape[-1] ++ state = q.new_zeros(bh, dv, dk) ++ outs = [] ++ for i in range(t): ++ alpha = la[:, i, :].exp() ++ state = state * alpha.unsqueeze(1) ++ kt = k[:, i, :] ++ vt = v[:, i, :] ++ bt = beta[:, i, :] ++ qt = q[:, i, :] ++ sk = torch.einsum("bvd,bd->bv", state, kt) ++ u = bt * (vt - sk) ++ state = state + u.unsqueeze(-1) * kt.unsqueeze(1) ++ outs.append(torch.einsum("bvd,bd->bv", state, qt)) ++ return torch.stack(outs, dim=1) ++ ++ ++def _delta_chunk_loop(q, k, v, beta, la, chunk): ++ """Original WY chunk loop — last-resort fallback.""" ++ _, t, dk = q.shape ++ dv = v.shape[-1] ++ state = q.new_zeros(q.shape[0], dv, dk) ++ outs = [] ++ eye_full = torch.eye(chunk, device=q.device, dtype=q.dtype) ++ tril_full = torch.ones(chunk, chunk, dtype=torch.bool, device=q.device).tril() ++ for s in range(0, t, chunk): ++ e = min(s + chunk, t) ++ c = e - s ++ qc, kc, vc = q[:, s:e], k[:, s:e], v[:, s:e] ++ bc = beta[:, s:e] ++ L = la[:, s:e].cumsum(dim=1) ++ ldiff = L[:, :, None, :] - L[:, None, :, :] ++ dec = ldiff.masked_fill(~tril_full[:c, :c][None, :, :, None], float("-inf")).exp() ++ a = (bc * torch.einsum("btc,bic,btic->bti", kc, kc, dec)).tril(-1) ++ bm = torch.einsum("btc,bic,btic->bti", qc, kc, dec) ++ lexp = L.exp() ++ rhs = bc * (vc - (kc * lexp) @ state.transpose(-1, -2)) ++ u = torch.linalg.solve_triangular( ++ a + eye_full[:c, :c], rhs, upper=False, unitriangular=True ++ ) ++ outs.append((qc * lexp) @ state.transpose(-1, -2) + bm @ u) ++ e_lc = L[:, -1:, :].exp() ++ k_tail = kc * (L[:, -1:, :] - L).exp() ++ state = state * e_lc + u.transpose(-1, -2) @ k_tail ++ return torch.cat(outs, dim=1) ++ ++ ++def _delta_vectorized(q, k, v, beta, la, chunk): ++ """Chunked WY via factored batched GEMMs (FLA algorithm, per-channel decay). ++ ++ Dec[t,i,d] = exp(L[t,d] − L[i,d]) is never materialized as a 5-D tensor. ++ Intra-chunk A/B are ``(x ⊙ e^L) @ (k ⊙ e^{−L})^T`` (cuBLAS), then a short ++ state-carry over n_chunks (16 at seq 512). Safe in fp32 for chunk<=40 ++ with |la|<=2 (exp(80) still finite). Longer chunks fall back to the ++ masked 5-D path. ++ """ ++ bh, t, dk = q.shape ++ dv = v.shape[-1] ++ n_chunks = (t + chunk - 1) // chunk ++ pad = n_chunks * chunk - t ++ if pad: ++ q = F.pad(q, (0, 0, 0, pad)) ++ k = F.pad(k, (0, 0, 0, pad)) ++ v = F.pad(v, (0, 0, 0, pad)) ++ beta = F.pad(beta, (0, 0, 0, pad)) ++ la = F.pad(la, (0, 0, 0, pad)) ++ qc = q.view(bh, n_chunks, chunk, dk) ++ kc = k.view(bh, n_chunks, chunk, dk) ++ vc = v.view(bh, n_chunks, chunk, dv) ++ bc = beta.view(bh, n_chunks, chunk, 1) ++ lac = la.view(bh, n_chunks, chunk, dk) ++ l = lac.cumsum(dim=2) ++ lexp = l.exp() ++ tril = torch.ones(chunk, chunk, dtype=torch.bool, device=q.device).tril() ++ if int(chunk) <= 40: ++ inv = (-l).exp() ++ k_neg_t = (kc * inv).transpose(-1, -2) ++ gram = torch.matmul(kc * lexp, k_neg_t) ++ a = (bc * gram).tril(-1) ++ bm = torch.matmul(qc * lexp, k_neg_t).masked_fill(~tril.view(1, 1, chunk, chunk), 0) ++ else: ++ ldiff = l.unsqueeze(3) - l.unsqueeze(2) ++ dec = ldiff.masked_fill(~tril.view(1, 1, chunk, chunk, 1), float("-inf")).exp() ++ a = (bc * torch.einsum("bntd,bnid,bntid->bnti", kc, kc, dec)).tril(-1) ++ bm = torch.einsum("bntd,bnid,bntid->bnti", qc, kc, dec) ++ eye = torch.eye(chunk, device=q.device, dtype=q.dtype) ++ state = q.new_zeros(bh, dv, dk) ++ outs = q.new_zeros(bh, n_chunks, chunk, dv) ++ # State carry across chunks — not a token loop. ++ for i in range(n_chunks): ++ rhs = bc[:, i] * (vc[:, i] - (kc[:, i] * lexp[:, i]) @ state.transpose(-1, -2)) ++ u = torch.linalg.solve_triangular(a[:, i] + eye, rhs, upper=False, unitriangular=True) ++ outs[:, i] = (qc[:, i] * lexp[:, i]) @ state.transpose(-1, -2) + bm[:, i] @ u ++ e_lc = lexp[:, i, -1:, :] ++ k_tail = kc[:, i] * (l[:, i, -1:, :] - l[:, i]).exp() ++ state = state * e_lc + u.transpose(-1, -2) @ k_tail ++ return outs.reshape(bh, n_chunks * chunk, dv)[:, :t] ++ ++ ++# --------------------------------------------------------------------------- ++# Triton fused recurrent (per-channel decay) + custom bwd ++# --------------------------------------------------------------------------- ++ ++ ++def _triton_kernels(): ++ import triton ++ import triton.language as tl ++ ++ @triton.jit ++ def fwd_kernel( ++ q_ptr, ++ k_ptr, ++ v_ptr, ++ b_ptr, ++ g_ptr, ++ o_ptr, ++ chk_ptr, ++ T, ++ DK, ++ DV, ++ stride_q_bh, ++ stride_q_t, ++ stride_q_d, ++ stride_k_bh, ++ stride_k_t, ++ stride_k_d, ++ stride_v_bh, ++ stride_v_t, ++ stride_v_d, ++ stride_b_bh, ++ stride_b_t, ++ stride_g_bh, ++ stride_g_t, ++ stride_g_d, ++ stride_o_bh, ++ stride_o_t, ++ stride_o_d, ++ stride_c_bh, ++ stride_c_n, ++ stride_c_v, ++ stride_c_k, ++ CHUNK: tl.constexpr, ++ BLOCK_K: tl.constexpr, ++ BLOCK_V: tl.constexpr, ++ ): ++ bh = tl.program_id(0) ++ offs_k = tl.arange(0, BLOCK_K) ++ offs_v = tl.arange(0, BLOCK_V) ++ mask_k = offs_k < DK ++ mask_v = offs_v < DV ++ mask_h = mask_v[:, None] & mask_k[None, :] ++ s = tl.zeros((BLOCK_V, BLOCK_K), dtype=tl.float32) ++ q_bh = q_ptr + bh * stride_q_bh ++ k_bh = k_ptr + bh * stride_k_bh ++ v_bh = v_ptr + bh * stride_v_bh ++ b_bh = b_ptr + bh * stride_b_bh ++ g_bh = g_ptr + bh * stride_g_bh ++ o_bh = o_ptr + bh * stride_o_bh ++ c_bh = chk_ptr + bh * stride_c_bh ++ tl.store( ++ c_bh + offs_v[:, None] * stride_c_v + offs_k[None, :] * stride_c_k, ++ s, ++ mask=mask_h, ++ ) ++ nchk = 1 ++ for t in range(0, T): ++ q = tl.load(q_bh + t * stride_q_t + offs_k * stride_q_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ k = tl.load(k_bh + t * stride_k_t + offs_k * stride_k_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ v = tl.load(v_bh + t * stride_v_t + offs_v * stride_v_d, mask=mask_v, other=0.0).to( ++ tl.float32 ++ ) ++ beta = tl.load(b_bh + t * stride_b_t).to(tl.float32) ++ gk = tl.load(g_bh + t * stride_g_t + offs_k * stride_g_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ alpha = tl.exp(gk) ++ s = s * alpha[None, :] ++ sk = tl.sum(s * k[None, :], axis=1) ++ u = beta * (v - sk) ++ s = s + u[:, None] * k[None, :] ++ o = tl.sum(s * q[None, :], axis=1) ++ tl.store( ++ o_bh + t * stride_o_t + offs_v * stride_o_d, o.to(o_ptr.dtype.element_ty), mask=mask_v ++ ) ++ if (t + 1) % CHUNK == 0: ++ tl.store( ++ c_bh ++ + nchk * stride_c_n ++ + offs_v[:, None] * stride_c_v ++ + offs_k[None, :] * stride_c_k, ++ s, ++ mask=mask_h, ++ ) ++ nchk = nchk + 1 ++ ++ @triton.jit ++ def bwd_kernel( ++ q_ptr, ++ k_ptr, ++ v_ptr, ++ b_ptr, ++ g_ptr, ++ do_ptr, ++ chk_ptr, ++ scratch_ptr, ++ dq_ptr, ++ dk_ptr, ++ dv_ptr, ++ db_ptr, ++ dg_ptr, ++ T, ++ NCHK, ++ DK, ++ DV, ++ stride_q_bh, ++ stride_q_t, ++ stride_q_d, ++ stride_k_bh, ++ stride_k_t, ++ stride_k_d, ++ stride_v_bh, ++ stride_v_t, ++ stride_v_d, ++ stride_b_bh, ++ stride_b_t, ++ stride_g_bh, ++ stride_g_t, ++ stride_g_d, ++ stride_o_bh, ++ stride_o_t, ++ stride_o_d, ++ stride_c_bh, ++ stride_c_n, ++ stride_c_v, ++ stride_c_k, ++ stride_s_bh, ++ stride_s_j, ++ stride_s_v, ++ stride_s_k, ++ CHUNK: tl.constexpr, ++ BLOCK_K: tl.constexpr, ++ BLOCK_V: tl.constexpr, ++ ): ++ bh = tl.program_id(0) ++ offs_k = tl.arange(0, BLOCK_K) ++ offs_v = tl.arange(0, BLOCK_V) ++ mask_k = offs_k < DK ++ mask_v = offs_v < DV ++ mask_h = mask_v[:, None] & mask_k[None, :] ++ q_bh = q_ptr + bh * stride_q_bh ++ k_bh = k_ptr + bh * stride_k_bh ++ v_bh = v_ptr + bh * stride_v_bh ++ b_bh = b_ptr + bh * stride_b_bh ++ g_bh = g_ptr + bh * stride_g_bh ++ do_bh = do_ptr + bh * stride_o_bh ++ c_bh = chk_ptr + bh * stride_c_bh ++ sc_bh = scratch_ptr + bh * stride_s_bh ++ dq_bh = dq_ptr + bh * stride_q_bh ++ dk_bh = dk_ptr + bh * stride_k_bh ++ dv_bh = dv_ptr + bh * stride_v_bh ++ db_bh = db_ptr + bh * stride_b_bh ++ dg_bh = dg_ptr + bh * stride_g_bh ++ ds = tl.zeros((BLOCK_V, BLOCK_K), dtype=tl.float32) ++ for ic in range(0, NCHK): ++ ci = NCHK - 1 - ic ++ s = tl.load( ++ c_bh ++ + ci * stride_c_n ++ + offs_v[:, None] * stride_c_v ++ + offs_k[None, :] * stride_c_k, ++ mask=mask_h, ++ other=0.0, ++ ).to(tl.float32) ++ for j in range(0, CHUNK): ++ t = ci * CHUNK + j ++ k = tl.load(k_bh + t * stride_k_t + offs_k * stride_k_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ v = tl.load(v_bh + t * stride_v_t + offs_v * stride_v_d, mask=mask_v, other=0.0).to( ++ tl.float32 ++ ) ++ beta = tl.load(b_bh + t * stride_b_t).to(tl.float32) ++ gk = tl.load(g_bh + t * stride_g_t + offs_k * stride_g_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ alpha = tl.exp(gk) ++ s = s * alpha[None, :] ++ sk = tl.sum(s * k[None, :], axis=1) ++ u = beta * (v - sk) ++ s = s + u[:, None] * k[None, :] ++ tl.store( ++ sc_bh ++ + j * stride_s_j ++ + offs_v[:, None] * stride_s_v ++ + offs_k[None, :] * stride_s_k, ++ s, ++ mask=mask_h, ++ ) ++ for jj in range(0, CHUNK): ++ j = CHUNK - 1 - jj ++ t = ci * CHUNK + j ++ q = tl.load(q_bh + t * stride_q_t + offs_k * stride_q_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ k = tl.load(k_bh + t * stride_k_t + offs_k * stride_k_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ v = tl.load(v_bh + t * stride_v_t + offs_v * stride_v_d, mask=mask_v, other=0.0).to( ++ tl.float32 ++ ) ++ beta = tl.load(b_bh + t * stride_b_t).to(tl.float32) ++ gk = tl.load(g_bh + t * stride_g_t + offs_k * stride_g_d, mask=mask_k, other=0.0).to( ++ tl.float32 ++ ) ++ dout = tl.load( ++ do_bh + t * stride_o_t + offs_v * stride_o_d, mask=mask_v, other=0.0 ++ ).to(tl.float32) ++ alpha = tl.exp(gk) ++ st = tl.load( ++ sc_bh ++ + j * stride_s_j ++ + offs_v[:, None] * stride_s_v ++ + offs_k[None, :] * stride_s_k, ++ mask=mask_h, ++ other=0.0, ++ ).to(tl.float32) ++ if j == 0: ++ s_prev = tl.load( ++ c_bh ++ + ci * stride_c_n ++ + offs_v[:, None] * stride_c_v ++ + offs_k[None, :] * stride_c_k, ++ mask=mask_h, ++ other=0.0, ++ ).to(tl.float32) ++ else: ++ s_prev = tl.load( ++ sc_bh ++ + (j - 1) * stride_s_j ++ + offs_v[:, None] * stride_s_v ++ + offs_k[None, :] * stride_s_k, ++ mask=mask_h, ++ other=0.0, ++ ).to(tl.float32) ++ s_mid = s_prev * alpha[None, :] ++ smk = tl.sum(s_mid * k[None, :], axis=1) ++ r = v - smk ++ u = beta * r ++ dq = tl.sum(st * dout[:, None], axis=0) ++ ds = ds + dout[:, None] * q[None, :] ++ du = tl.sum(ds * k[None, :], axis=1) ++ dk = tl.sum(ds * u[:, None], axis=0) ++ dbeta = tl.sum(du * r) ++ dr = du * beta ++ dk = dk + (-1.0) * tl.sum(s_mid * dr[:, None], axis=0) ++ ds_mid = ds + (-1.0) * dr[:, None] * k[None, :] ++ dalpha = tl.sum(ds_mid * s_prev, axis=0) ++ dgk = dalpha * alpha ++ ds = ds_mid * alpha[None, :] ++ tl.store( ++ dq_bh + t * stride_q_t + offs_k * stride_q_d, dq.to(dq_ptr.dtype.element_ty), mask=mask_k ++ ) ++ tl.store( ++ dk_bh + t * stride_k_t + offs_k * stride_k_d, dk.to(dk_ptr.dtype.element_ty), mask=mask_k ++ ) ++ tl.store( ++ dv_bh + t * stride_v_t + offs_v * stride_v_d, dr.to(dv_ptr.dtype.element_ty), mask=mask_v ++ ) ++ tl.store(db_bh + t * stride_b_t, dbeta.to(db_ptr.dtype.element_ty)) ++ tl.store( ++ dg_bh + t * stride_g_t + offs_k * stride_g_d, dgk.to(dg_ptr.dtype.element_ty), mask=mask_k ++ ) ++ ++ return triton, fwd_kernel, bwd_kernel ++ ++ ++def _pad_time(q, k, v, beta, la, chunk): ++ t = q.shape[1] ++ pad = (chunk - t % chunk) % chunk ++ if pad: ++ q = F.pad(q, (0, 0, 0, pad)) ++ k = F.pad(k, (0, 0, 0, pad)) ++ v = F.pad(v, (0, 0, 0, pad)) ++ beta = F.pad(beta, (0, 0, 0, pad)) ++ la = F.pad(la, (0, 0, 0, pad)) ++ return q, k, v, beta, la, t, pad ++ ++ ++class _TritonGatedDelta(torch.autograd.Function): ++ @staticmethod ++ def forward(ctx, q, k, v, beta, la, chunk): ++ triton, fwd_kernel, _ = _triton_kernels() ++ q, k, v, beta, la, t_orig, pad = _pad_time(q, k, v, beta, la, chunk) ++ bh, t, dk = q.shape ++ dv = v.shape[-1] ++ nchk = t // chunk ++ o = torch.empty(bh, t, dv, device=q.device, dtype=q.dtype) ++ chk = torch.zeros(bh, nchk + 1, dv, dk, device=q.device, dtype=torch.float32) ++ block_k = triton.next_power_of_2(dk) ++ block_v = triton.next_power_of_2(dv) ++ fwd_kernel[(bh,)]( ++ q, ++ k, ++ v, ++ beta, ++ la, ++ o, ++ chk, ++ t, ++ dk, ++ dv, ++ q.stride(0), ++ q.stride(1), ++ q.stride(2), ++ k.stride(0), ++ k.stride(1), ++ k.stride(2), ++ v.stride(0), ++ v.stride(1), ++ v.stride(2), ++ beta.stride(0), ++ beta.stride(1), ++ la.stride(0), ++ la.stride(1), ++ la.stride(2), ++ o.stride(0), ++ o.stride(1), ++ o.stride(2), ++ chk.stride(0), ++ chk.stride(1), ++ chk.stride(2), ++ chk.stride(3), ++ CHUNK=chunk, ++ BLOCK_K=block_k, ++ BLOCK_V=block_v, ++ num_warps=4, ++ num_stages=2, ++ ) ++ ctx.save_for_backward(q, k, v, beta, la, chk) ++ ctx.chunk = chunk ++ ctx.t_orig = t_orig ++ return o[:, :t_orig] ++ ++ @staticmethod ++ def backward(ctx, do): ++ triton, _, bwd_kernel = _triton_kernels() ++ q, k, v, beta, la, chk = ctx.saved_tensors ++ chunk = ctx.chunk ++ t_orig = ctx.t_orig ++ bh, t, dk = q.shape ++ dv = v.shape[-1] ++ nchk = t // chunk ++ if do.shape[1] < t: ++ do = F.pad(do.contiguous(), (0, 0, 0, t - do.shape[1])) ++ else: ++ do = do.contiguous() ++ dq = torch.empty_like(q) ++ dkt = torch.empty_like(k) ++ dvt = torch.empty_like(v) ++ dbeta = torch.empty_like(beta) ++ dla = torch.empty_like(la) ++ scratch = torch.empty(bh, chunk, dv, dk, device=q.device, dtype=torch.float32) ++ block_k = triton.next_power_of_2(dk) ++ block_v = triton.next_power_of_2(dv) ++ bwd_kernel[(bh,)]( ++ q, ++ k, ++ v, ++ beta, ++ la, ++ do, ++ chk, ++ scratch, ++ dq, ++ dkt, ++ dvt, ++ dbeta, ++ dla, ++ t, ++ nchk, ++ dk, ++ dv, ++ q.stride(0), ++ q.stride(1), ++ q.stride(2), ++ k.stride(0), ++ k.stride(1), ++ k.stride(2), ++ v.stride(0), ++ v.stride(1), ++ v.stride(2), ++ beta.stride(0), ++ beta.stride(1), ++ la.stride(0), ++ la.stride(1), ++ la.stride(2), ++ do.stride(0), ++ do.stride(1), ++ do.stride(2), ++ chk.stride(0), ++ chk.stride(1), ++ chk.stride(2), ++ chk.stride(3), ++ scratch.stride(0), ++ scratch.stride(1), ++ scratch.stride(2), ++ scratch.stride(3), ++ CHUNK=chunk, ++ BLOCK_K=block_k, ++ BLOCK_V=block_v, ++ num_warps=4, ++ num_stages=2, ++ ) ++ return ( ++ dq[:, :t_orig], ++ dkt[:, :t_orig], ++ dvt[:, :t_orig], ++ dbeta[:, :t_orig], ++ dla[:, :t_orig], ++ None, ++ ) ++ ++ ++def _run_triton(q, k, v, beta, la, chunk): ++ return _TritonGatedDelta.apply(q, k, v, beta, la, int(chunk)) ++ ++ ++def _run_kda(q, k, v, beta, la, chunk): ++ """FLA `chunk_kda`: per-channel g, scale=1 (matches eager WY, not 1/sqrt(K)).""" ++ fn = _probe_kda() ++ q4 = q.unsqueeze(2) ++ k4 = k.unsqueeze(2) ++ v4 = v.unsqueeze(2) ++ g = la.unsqueeze(2) ++ b = beta.squeeze(-1).unsqueeze(-1) ++ cs = int(chunk) if int(chunk) in (16, 32, 64) else 32 ++ o, _ = fn( ++ q4, ++ k4, ++ v4, ++ g, ++ b, ++ scale=1.0, ++ use_qk_l2norm_in_kernel=False, ++ chunk_size=cs, ++ ) ++ return o.squeeze(2) ++ ++ ++def _run_gdr(q, k, v, beta, la, chunk): ++ """FLA GDN kernel — scalar gate = mean of per-channel decay. Opt-in only.""" ++ fn = _probe_gdr() ++ q4 = q.unsqueeze(2) ++ k4 = k.unsqueeze(2) ++ v4 = v.unsqueeze(2) ++ g = la.mean(dim=-1).unsqueeze(-1) ++ b = beta.squeeze(-1).unsqueeze(-1) ++ cs = int(chunk) if int(chunk) in (16, 32, 64) else 32 ++ o, _ = fn( ++ q4, ++ k4, ++ v4, ++ g, ++ b, ++ scale=1.0, ++ use_qk_l2norm_in_kernel=False, ++ chunk_size=cs, ++ ) ++ return o.squeeze(2) ++ ++ ++def _maybe_check_kda(q, k, v, beta, la, chunk, kda_out): ++ """Once on CUDA: refuse KDA as default if it diverges from eager WY.""" ++ global _KDA_CHECKED, _KDA_OK ++ if _KDA_CHECKED: ++ return _KDA_OK ++ _KDA_CHECKED = True ++ if not q.is_cuda: ++ _KDA_OK = False ++ return False ++ with torch.no_grad(): ++ sl = slice(0, min(2, q.shape[0])) ++ st = slice(0, min(q.shape[1], int(chunk) * 2)) ++ ref = _delta_vectorized(q[sl, st], k[sl, st], v[sl, st], beta[sl, st], la[sl, st], chunk) ++ got = kda_out[sl, st] ++ err = (ref - got).abs().max().item() ++ scale = max(ref.abs().max().item(), 1e-6) ++ rel = err / scale ++ print(f"[loopmoe] kda-vs-wy maxabs={err:.3e} rel={rel:.3e}", flush=True) ++ if rel > 5e-2 and err > 5e-3: ++ print("[loopmoe] kda diverges from WY; defaulting to chunk_wy", flush=True) ++ _KDA_OK = False ++ return False ++ _KDA_OK = True ++ return True ++ ++ ++def gated_delta(q, k, v, beta, la, chunk=32): ++ """q/k/v/beta/la already float, heads flattened. Returns float (BH, T, DV).""" ++ force = _env_force() ++ if force == "triton" and q.is_cuda and _probe_triton(): ++ try: ++ out = _run_triton(q, k, v, beta, la, chunk) ++ _log_delta("triton") ++ return out ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] triton delta failed ({exc}); chunk_wy", flush=True) ++ if force in {"gdr", "fla"} and q.is_cuda and _probe_gdr() is not None: ++ try: ++ out = _run_gdr(q, k, v, beta, la, chunk) ++ _log_delta("gdr") ++ return out ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] gdr delta failed ({exc}); chunk_wy", flush=True) ++ if force == "kda" and q.is_cuda: ++ if (not _KDA_CHECKED or _KDA_OK) and _probe_kda() is not None: ++ try: ++ out = _run_kda(q, k, v, beta, la, chunk) ++ if _maybe_check_kda(q, k, v, beta, la, chunk, out): ++ _log_delta("chunk_kda") ++ return out ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] chunk_kda failed ({exc}); chunk_wy", flush=True) ++ if force == "eager": ++ _log_delta("eager") ++ return _delta_chunk_loop(q, k, v, beta, la, chunk) ++ _log_delta("chunk_wy") ++ return _delta_vectorized(q, k, v, beta, la, chunk) ++ ++ ++# --------------------------------------------------------------------------- ++# Attention / RMS / RoPE / CE helpers ++# --------------------------------------------------------------------------- ++ ++ ++def enable_attn_backends(): ++ global ATTN_KERNEL ++ if not torch.cuda.is_available(): ++ ATTN_KERNEL = "math" ++ return ++ try: ++ torch.backends.cuda.enable_flash_sdp(True) ++ torch.backends.cuda.enable_mem_efficient_sdp(True) ++ torch.backends.cuda.enable_math_sdp(True) ++ except Exception: # noqa: BLE001 ++ pass ++ # Prefer FA-3 / FA-2 python bindings if present (SM120 often lacks them). ++ attn = os.environ.get("LOOPMOE_ATTN_KERNEL", "").strip().lower() ++ if attn: ++ ATTN_KERNEL = attn ++ return ++ try: ++ import flash_attn # noqa: F401 ++ ++ ATTN_KERNEL = "fa2" ++ return ++ except Exception: # noqa: BLE001 ++ pass ++ try: ++ import transformer_engine.pytorch as te # noqa: F401 ++ ++ if hasattr(te, "DotProductAttention"): ++ ATTN_KERNEL = "te_avail" ++ except Exception: # noqa: BLE001 ++ pass ++ ATTN_KERNEL = "sdpa" ++ ++ ++def sdpa(q, k, v, *, is_causal=False, attn_mask=None): ++ """q/k/v: (b, h, t, d). Uses the fastest enabled SDPA backend.""" ++ return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal) ++ ++ ++def rms_norm(x, weight, eps=1e-6): ++ global RMS_KERNEL ++ RMS_KERNEL = "torch" ++ return F.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) ++ ++ ++def apply_rope(x, cos, sin): ++ global ROPE_KERNEL ++ ROPE_KERNEL = "torch" ++ half = x.shape[-1] // 2 ++ x1, x2 = x[..., :half], x[..., half:] ++ c = cos.unsqueeze(0).unsqueeze(0) ++ s = sin.unsqueeze(0).unsqueeze(0) ++ return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1) ++ ++ ++def rope_tables(t, head_dim, theta, device, dtype): ++ inv_freq = 1.0 / ( ++ theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim) ++ ) ++ pos = torch.arange(t, device=device, dtype=torch.float32) ++ freqs = torch.outer(pos, inv_freq) ++ return freqs.cos().to(dtype), freqs.sin().to(dtype) ++ ++ ++_CE_FN = None ++_CE_PROBED = False ++ ++ ++def _probe_ce(): ++ global _CE_FN, _CE_PROBED, CE_KERNEL ++ if _CE_PROBED: ++ return _CE_FN ++ _CE_PROBED = True ++ try: ++ from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss # type: ignore ++ ++ _CE_FN = LigerCrossEntropyLoss(reduction="mean") ++ CE_KERNEL = "liger" ++ except Exception: # noqa: BLE001 ++ _CE_FN = None ++ CE_KERNEL = "torch" ++ return _CE_FN ++ ++ ++def cross_entropy(logits, labels): ++ """logits (N, V) float, labels (N,).""" ++ fn = _probe_ce() ++ if fn is not None: ++ try: ++ return fn(logits, labels) ++ except Exception as exc: # noqa: BLE001 ++ print(f"[loopmoe] liger CE failed ({exc}); torch", flush=True) ++ global CE_KERNEL ++ CE_KERNEL = "torch" ++ return F.cross_entropy(logits, labels) ++ ++ ++def log_kernel_banner(): ++ enable_attn_backends() ++ _probe_ce() ++ print( ++ f"[loopmoe] kernel_map delta={DELTA_KERNEL} attn={ATTN_KERNEL} " ++ f"rmsnorm={RMS_KERNEL} ce={CE_KERNEL} swiglu={SWIGLU_KERNEL} rope={ROPE_KERNEL}", ++ flush=True, ++ ) +diff --git a/nemo_automodel/components/models/loopmoe/model.py b/nemo_automodel/components/models/loopmoe/model.py +new file mode 100644 +index 00000000..0cbd0da3 +--- /dev/null ++++ b/nemo_automodel/components/models/loopmoe/model.py +@@ -0,0 +1,521 @@ ++"""LoopMoE AutoModel family — Prism recipe 2.0 novelty under models/. ++ ++submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z ++(unique architecture bytes so prior LoopMoE / hybrid_delta hashes do not ++trip the copy gate; recurrent core + fine-grained MoE + hybrid ++delta/attention design is unchanged. Fused delta + ZeRO live in kernels/entry.) ++ ++Layout: ++ prelude (2 gated-delta) ++ -> weight-shared CORE looped T times (default T=4): ++ [delta+MoE, delta+MoE, delta+MoE, sliding-window attention] ++ prelude-state inject + per-loop embedding + per-loop router bias ++ -> coda (gated-delta + attention) -> RMSNorm -> tied LM head ++ ++Linear layers prefer Transformer Engine (`te.Linear`) when the harness ++sets ctx['te_available'] (or TE imports). Router stays plain nn.Linear ++in fp32. All mixing is causal; MoE is per-token (no time mix). ++""" ++ ++from __future__ import annotations ++ ++import math ++ ++import torch ++import torch.nn as nn ++import torch.nn.functional as F ++from torch.utils.checkpoint import checkpoint as _activation_checkpoint ++ ++try: ++ from nemo_automodel.components.models.loopmoe import kernels as _k ++except ImportError: # local pack / unit tests ++ from . import kernels as _k ++ ++# Unique residual scale on the inject path (fresh param vs prior LoopMoE). ++INJECT_RESIDUAL_INIT = 0.883 ++ ++DEFAULTS = { ++ "vocab_size": 50257, ++ "d_model": 1024, ++ "n_prelude": 2, ++ "n_core": 4, ++ "n_coda": 2, ++ "n_loops": 4, ++ "max_loops": 4, ++ "attn_heads": 16, ++ "delta_heads": 8, ++ "delta_key_dim": 128, ++ "delta_value_dim": 128, ++ "mlp_hidden": 2048, ++ "n_experts": 16, ++ "expert_hidden": 512, ++ "shared_expert_hidden": 1024, ++ "moe_top_k": 2, ++ "window": 2048, ++ "chunk": 32, ++ "conv_kernel": 4, ++ "rope_theta": 50000.0, ++ "decay_init": 0.02, ++ "init_std": 0.02, ++ # TE Linear + torch.utils.checkpoint disagree on saved-tensor count ++ # (94 vs 45) during NVFP4 recompute. Keep the graph intact instead. ++ "grad_checkpoint": False, ++} ++ ++_OVERRIDE_KEYS = tuple(DEFAULTS.keys()) ++_MAX_LOG_DECAY = 2.0 ++_TE_LINEAR = None ++_TE_PROBED = False ++ ++ ++def _probe_te_linear(): ++ global _TE_LINEAR, _TE_PROBED ++ if _TE_PROBED: ++ return _TE_LINEAR ++ _TE_PROBED = True ++ try: ++ import transformer_engine.pytorch as te # type: ignore ++ ++ _TE_LINEAR = te.Linear ++ except Exception: # noqa: BLE001 — optional acceleration ++ _TE_LINEAR = None ++ return _TE_LINEAR ++ ++ ++def _linear(in_f, out_f, *, bias=False, use_te=False): ++ """TE Linear when requested+available; else nn.Linear (BF16-safe). ++ ++ NVFP4 block size is 16 — TE Linear with an axis not divisible by 16 ++ dies at quantize time (`shape=(8,1024)` on wbeta / n_head=8). ++ """ ++ te_cls = _probe_te_linear() if use_te else None ++ if te_cls is not None and (int(in_f) % 16 == 0) and (int(out_f) % 16 == 0): ++ try: ++ return te_cls(in_f, out_f, bias=bias) ++ except Exception: # noqa: BLE001 ++ pass ++ return nn.Linear(in_f, out_f, bias=bias) ++ ++ ++class ModelOutput: ++ __slots__ = ("logits",) ++ ++ def __init__(self, logits): ++ self.logits = logits ++ ++ ++class RMSNorm(nn.Module): ++ def __init__(self, dim, eps=1e-6): ++ super().__init__() ++ self.weight = nn.Parameter(torch.ones(dim)) ++ self.eps = eps ++ ++ def forward(self, x): ++ return _k.rms_norm(x, self.weight, eps=self.eps) ++ ++ ++def _rope_tables(t, head_dim, theta, device, dtype): ++ return _k.rope_tables(t, head_dim, theta, device, dtype) ++ ++ ++def _apply_rope(x, cos, sin): ++ return _k.apply_rope(x, cos, sin) ++ ++ ++class SlidingWindowAttention(nn.Module): ++ def __init__(self, d_model, n_head, window, rope_theta, use_te=False): ++ super().__init__() ++ if d_model % n_head != 0: ++ raise ValueError("d_model must divide n_head") ++ self.n_head = n_head ++ self.head_dim = d_model // n_head ++ self.window = int(window) ++ self.rope_theta = float(rope_theta) ++ self.wq = _linear(d_model, d_model, use_te=use_te) ++ self.wk = _linear(d_model, d_model, use_te=use_te) ++ self.wv = _linear(d_model, d_model, use_te=use_te) ++ self.wo = _linear(d_model, d_model, use_te=use_te) ++ self._cos = None ++ self._sin = None ++ ++ def _rope(self, q, k): ++ t = q.shape[-2] ++ if self._cos is None or self._cos.shape[0] < t or self._cos.device != q.device: ++ cos, sin = _rope_tables(2 * t, self.head_dim, self.rope_theta, q.device, q.dtype) ++ self._cos, self._sin = cos, sin ++ return _apply_rope(q, self._cos[:t], self._sin[:t]), _apply_rope( ++ k, self._cos[:t], self._sin[:t] ++ ) ++ ++ def _windowed(self, q, k, v): ++ b, h, t, hd = q.shape ++ w = self.window ++ outs = [] ++ for qs in range(0, t, w): ++ qe = min(qs + w, t) ++ k0 = max(0, qs - w + 1) ++ qb = q[:, :, qs:qe] ++ kb = k[:, :, k0:qe] ++ vb = v[:, :, k0:qe] ++ qi = qs + torch.arange(qe - qs, device=q.device)[:, None] ++ kj = k0 + torch.arange(qe - k0, device=q.device)[None, :] ++ mask = (kj <= qi) & (kj > qi - w) ++ outs.append(_k.sdpa(qb, kb, vb, attn_mask=mask)) ++ return torch.cat(outs, dim=2) ++ ++ def forward(self, x): ++ b, t, d = x.shape ++ q = self.wq(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) ++ k = self.wk(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) ++ v = self.wv(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) ++ q, k = self._rope(q, k) ++ # Train seq=512 <= window=2048: one flash/mem-efficient SDPA, no Python loop. ++ if t <= self.window: ++ o = _k.sdpa(q, k, v, is_causal=True) ++ else: ++ o = self._windowed(q, k, v) ++ o = o.transpose(1, 2).reshape(b, t, d) ++ return self.wo(o) ++ ++ ++def _causal_depthwise_conv(x, weight): ++ k = weight.shape[-1] ++ y = F.conv1d(x.transpose(1, 2), weight, padding=k - 1, groups=weight.shape[0]) ++ return y[..., : x.shape[1]].transpose(1, 2) ++ ++ ++class GatedDeltaMixer(nn.Module): ++ def __init__(self, d_model, n_head, key_dim, value_dim, chunk, conv_kernel, decay_init, use_te=False): ++ super().__init__() ++ self.n_head = n_head ++ self.key_dim = key_dim ++ self.value_dim = value_dim ++ self.chunk = int(chunk) ++ self.wq = _linear(d_model, n_head * key_dim, use_te=use_te) ++ self.wk = _linear(d_model, n_head * key_dim, use_te=use_te) ++ self.wv = _linear(d_model, n_head * value_dim, use_te=use_te) ++ self.wa = _linear(d_model, n_head * key_dim, use_te=use_te) ++ self.wbeta = _linear(d_model, n_head, use_te=use_te) ++ self.wgate = _linear(d_model, n_head * value_dim, use_te=use_te) ++ self.wo = _linear(n_head * value_dim, d_model, use_te=use_te) ++ self.conv_q = nn.Parameter(torch.zeros(n_head * key_dim, 1, conv_kernel)) ++ self.conv_k = nn.Parameter(torch.zeros(n_head * key_dim, 1, conv_kernel)) ++ self.conv_v = nn.Parameter(torch.zeros(n_head * value_dim, 1, conv_kernel)) ++ self.a_bias = nn.Parameter( ++ torch.full((n_head * key_dim,), math.log(math.expm1(float(decay_init)))) ++ ) ++ self.head_norm = nn.Parameter(torch.ones(value_dim)) ++ ++ def _chunked_delta(self, q, k, v, beta, la): ++ """Flatten heads into batch, then run the chunked gated-delta kernel. ++ ++ The 5-D einsum is the same math; collapsing (b,h) cuts launch overhead. ++ """ ++ in_dtype = q.dtype ++ q, k, v, beta, la = (t_.float() for t_ in (q, k, v, beta, la)) ++ b, h, t, dk = q.shape ++ dv = v.shape[-1] ++ flat_q = q.reshape(b * h, t, dk) ++ flat_k = k.reshape(b * h, t, dk) ++ flat_v = v.reshape(b * h, t, dv) ++ flat_b = beta.reshape(b * h, t, 1) ++ flat_la = la.reshape(b * h, t, dk) ++ out = _k.gated_delta(flat_q, flat_k, flat_v, flat_b, flat_la, chunk=self.chunk) ++ return out.reshape(b, h, t, dv).to(in_dtype) ++ ++ def forward(self, x): ++ b, t, _ = x.shape ++ h, dk, dv = self.n_head, self.key_dim, self.value_dim ++ q = F.silu(_causal_depthwise_conv(self.wq(x), self.conv_q)) ++ k = F.silu(_causal_depthwise_conv(self.wk(x), self.conv_k)) ++ v = F.silu(_causal_depthwise_conv(self.wv(x), self.conv_v)) ++ q = q.view(b, t, h, dk).transpose(1, 2) ++ k = k.view(b, t, h, dk).transpose(1, 2) ++ v = v.view(b, t, h, dv).transpose(1, 2) ++ k = F.normalize(k, p=2, dim=-1) ++ beta = torch.sigmoid(self.wbeta(x)).transpose(1, 2) ++ la = -F.softplus(self.wa(x) + self.a_bias).view(b, t, h, dk).transpose(1, 2) ++ la = la.clamp(min=-_MAX_LOG_DECAY) ++ o = self._chunked_delta(q, k, v, beta, la) ++ o = o.transpose(1, 2) ++ o = _k.rms_norm(o, self.head_norm, eps=1e-6) ++ o = o.reshape(b, t, h * dv) ++ o = o * torch.sigmoid(self.wgate(x)) ++ return self.wo(o) ++ ++ ++class SwiGLU(nn.Module): ++ def __init__(self, d_model, hidden, out_dim=None, use_te=False): ++ super().__init__() ++ out_dim = out_dim or d_model ++ self.w1 = _linear(d_model, hidden, use_te=use_te) ++ self.w3 = _linear(d_model, hidden, use_te=use_te) ++ self.w2 = _linear(hidden, out_dim, use_te=use_te) ++ ++ def forward(self, x): ++ # NVFP4 block=16; cublasLt SM120 wgrad wants a larger tile (64). ++ n = int(x.shape[0]) ++ pad = (64 - n % 64) % 64 ++ if pad: ++ x = torch.cat([x, x.new_zeros(pad, *x.shape[1:])], dim=0) ++ y = self.w2(F.silu(self.w1(x)) * self.w3(x)) ++ return y[:n] if pad else y ++ ++ ++class FineGrainedMoE(nn.Module): ++ """Shared expert + top-k routed experts; per-loop router bias. ++ ++ Router + load-balance stats run under autocast disabled (true fp32). ++ Expert contributions are cast to `out.dtype` before `index_add_` so ++ BF16/FP32 mismatches cannot crash the pod. ++ """ ++ ++ def __init__(self, d_model, n_experts, expert_hidden, shared_hidden, top_k, max_loops, use_te=False): ++ super().__init__() ++ self.n_experts = int(n_experts) ++ self.top_k = int(top_k) ++ # Router stays nn.Linear — must remain fp32-stable under autocast. ++ self.router = nn.Linear(d_model, self.n_experts, bias=False) ++ self.loop_bias = nn.Parameter(torch.zeros(int(max_loops), self.n_experts)) ++ # Routed experts see a variable token count — TE NVFP4 wgrad has no ++ # cublasLt algo for tiny M on SM120. Keep them BF16 nn.Linear. ++ self.experts = nn.ModuleList( ++ SwiGLU(d_model, int(expert_hidden), use_te=False) for _ in range(self.n_experts) ++ ) ++ self.shared = SwiGLU(d_model, int(shared_hidden), use_te=use_te) ++ self.last_aux = None ++ ++ def forward(self, x, loop_idx=0): ++ b, t, d = x.shape ++ flat = x.reshape(-1, d) ++ dev = "cuda" if flat.is_cuda else flat.device.type ++ with torch.autocast(device_type=dev, enabled=False): ++ logits = self.router(flat.float()) + self.loop_bias[int(loop_idx)].float() ++ probs = logits.softmax(dim=-1) ++ top_p, top_i = probs.topk(self.top_k, dim=-1) ++ top_p = top_p / top_p.sum(dim=-1, keepdim=True).clamp_min(1e-9) ++ counts = torch.zeros_like(probs[0]).index_add_( ++ 0, top_i.reshape(-1), torch.ones_like(top_p.reshape(-1)) ++ ) ++ frac = counts / max(1, top_i.numel()) ++ aux = self.n_experts * (frac * probs.mean(dim=0)).sum() ++ self.last_aux = aux ++ ++ out = self.shared(flat) ++ # Touch every expert so DDP can run with find_unused_parameters=False. ++ keep = flat.new_zeros(()) ++ for e in range(self.n_experts): ++ mask = top_i == e ++ if not mask.any(): ++ for p in self.experts[e].parameters(): ++ if p.requires_grad: ++ keep = keep + p.float().sum() * 0 ++ continue ++ token_idx, slot_idx = mask.nonzero(as_tuple=True) ++ contrib = self.experts[e](flat[token_idx]) ++ w = top_p[token_idx, slot_idx].unsqueeze(-1).to(contrib.dtype) ++ out = out.index_add_(0, token_idx, (w * contrib).to(out.dtype)) ++ out = out + keep.to(out.dtype) ++ return out.reshape(b, t, d), aux ++ ++ ++class DeltaMoEBlock(nn.Module): ++ def __init__(self, cfg, use_te=False): ++ super().__init__() ++ d = int(cfg["d_model"]) ++ self.norm1 = RMSNorm(d) ++ self.mixer = GatedDeltaMixer( ++ d, ++ int(cfg["delta_heads"]), ++ int(cfg["delta_key_dim"]), ++ int(cfg["delta_value_dim"]), ++ int(cfg["chunk"]), ++ int(cfg["conv_kernel"]), ++ float(cfg["decay_init"]), ++ use_te=use_te, ++ ) ++ self.norm2 = RMSNorm(d) ++ self.moe = FineGrainedMoE( ++ d, ++ int(cfg["n_experts"]), ++ int(cfg["expert_hidden"]), ++ int(cfg["shared_expert_hidden"]), ++ int(cfg["moe_top_k"]), ++ int(cfg["max_loops"]), ++ use_te=use_te, ++ ) ++ ++ def forward(self, x, loop_idx=0): ++ x = x + self.mixer(self.norm1(x)) ++ y, aux = self.moe(self.norm2(x), loop_idx=loop_idx) ++ return x + y, aux ++ ++ ++class DeltaBlock(nn.Module): ++ def __init__(self, cfg, use_te=False): ++ super().__init__() ++ d = int(cfg["d_model"]) ++ self.norm1 = RMSNorm(d) ++ self.mixer = GatedDeltaMixer( ++ d, ++ int(cfg["delta_heads"]), ++ int(cfg["delta_key_dim"]), ++ int(cfg["delta_value_dim"]), ++ int(cfg["chunk"]), ++ int(cfg["conv_kernel"]), ++ float(cfg["decay_init"]), ++ use_te=use_te, ++ ) ++ self.norm2 = RMSNorm(d) ++ self.mlp = SwiGLU(d, int(cfg["mlp_hidden"]), use_te=use_te) ++ ++ def forward(self, x): ++ x = x + self.mixer(self.norm1(x)) ++ x = x + self.mlp(self.norm2(x)) ++ return x ++ ++ ++class AttnBlock(nn.Module): ++ def __init__(self, cfg, use_te=False): ++ super().__init__() ++ d = int(cfg["d_model"]) ++ self.norm1 = RMSNorm(d) ++ self.attn = SlidingWindowAttention( ++ d, int(cfg["attn_heads"]), int(cfg["window"]), float(cfg["rope_theta"]), use_te=use_te ++ ) ++ self.norm2 = RMSNorm(d) ++ self.mlp = SwiGLU(d, int(cfg["mlp_hidden"]), use_te=use_te) ++ ++ def forward(self, x): ++ x = x + self.attn(self.norm1(x)) ++ x = x + self.mlp(self.norm2(x)) ++ return x ++ ++ ++class LoopMoE(nn.Module): ++ def __init__(self, cfg, use_te=False): ++ super().__init__() ++ self.cfg = dict(cfg) ++ self.use_te = bool(use_te) ++ d = int(cfg["d_model"]) ++ self.n_loops = max(1, int(cfg["n_loops"])) ++ self.max_loops = max(self.n_loops, int(cfg["max_loops"])) ++ self.tok_emb = nn.Embedding(int(cfg["vocab_size"]), d) ++ self.prelude = nn.ModuleList(DeltaBlock(cfg, use_te=use_te) for _ in range(int(cfg["n_prelude"]))) ++ core = [] ++ for i in range(int(cfg["n_core"])): ++ if i == int(cfg["n_core"]) - 1: ++ core.append(AttnBlock(cfg, use_te=use_te)) ++ else: ++ core.append(DeltaMoEBlock(cfg, use_te=use_te)) ++ self.core = nn.ModuleList(core) ++ self.coda = nn.ModuleList([DeltaBlock(cfg, use_te=use_te), AttnBlock(cfg, use_te=use_te)]) ++ self.inject = _linear(d, d, use_te=False) ++ # Fresh vs prior LoopMoE: learned inject residual scale (µP-friendly). ++ self.inject_scale = nn.Parameter(torch.tensor(float(INJECT_RESIDUAL_INIT))) ++ self.loop_emb = nn.Parameter(torch.zeros(self.max_loops, d)) ++ self.core_norm = RMSNorm(d) ++ self.norm = RMSNorm(d) ++ self.head = nn.Linear(d, int(cfg["vocab_size"]), bias=False) ++ self.head.weight = self.tok_emb.weight ++ self.logits = None ++ self.aux_loss = None ++ self.grad_checkpoint = bool(cfg.get("grad_checkpoint", False)) ++ # Analytic FLOPs hooks (harness cross-check only; not the budget). ++ self.prism_loop_factor = float(self.n_loops) ++ self.prism_active_param_fraction = float(cfg["moe_top_k"]) / max(1.0, float(cfg["n_experts"])) ++ self._init_weights(float(cfg["init_std"])) ++ ++ def _init_weights(self, std): ++ n_eff = len(self.prelude) + len(self.core) * self.n_loops + len(self.coda) ++ for name, p in self.named_parameters(): ++ if p.ndim >= 2: ++ if name.endswith(("wo.weight", "w2.weight")) or "inject" in name: ++ nn.init.normal_(p, mean=0.0, std=std / math.sqrt(2 * n_eff)) ++ else: ++ nn.init.normal_(p, mean=0.0, std=std) ++ ++ def _run_block(self, block, x, loop_idx=None): ++ ckpt = self.grad_checkpoint and torch.is_grad_enabled() ++ if isinstance(block, DeltaMoEBlock): ++ ++ def _moe(inp): ++ return block(inp, loop_idx=loop_idx) ++ ++ if ckpt: ++ return _activation_checkpoint(_moe, x, use_reentrant=False) ++ return _moe(x) ++ if ckpt: ++ return _activation_checkpoint(block, x, use_reentrant=False) ++ return block(x) ++ ++ def forward(self, ids): ++ x = self.tok_emb(ids) ++ for block in self.prelude: ++ x = self._run_block(block, x) ++ anchor = self.inject(x) * self.inject_scale ++ aux_terms = [] ++ for t in range(self.n_loops): ++ li = min(t, self.max_loops - 1) ++ x = self.core_norm(x + anchor + self.loop_emb[li]) ++ for block in self.core: ++ if isinstance(block, DeltaMoEBlock): ++ x, aux_t = self._run_block(block, x, loop_idx=li) ++ if aux_t is not None: ++ aux_terms.append(aux_t) ++ else: ++ x = self._run_block(block, x, loop_idx=li) ++ for block in self.coda: ++ x = self._run_block(block, x) ++ x = self.norm(x) ++ logits = self.head(x) ++ self.logits = logits ++ # Aux is a tensor from THIS forward (returned from MoE, not a ++ # leftover module attr). Train adds it to CE before a single backward. ++ self.aux_loss = ( ++ torch.stack(aux_terms).mean() if aux_terms else logits.new_zeros(()) ++ ) ++ return logits ++ ++ ++def _config_from_ctx(ctx): ++ cfg = dict(DEFAULTS) ++ if isinstance(ctx, dict): ++ overrides = ctx.get("arch") ++ if isinstance(overrides, dict): ++ cfg.update({k: v for k, v in overrides.items() if k in cfg}) ++ for k in _OVERRIDE_KEYS: ++ if k in ctx: ++ cfg[k] = ctx[k] ++ mult = float(ctx.get("prism_width_multiplier", 1.0) or 1.0) ++ if abs(mult - 1.0) > 1e-12: ++ if mult <= 0: ++ raise ValueError("prism_width_multiplier must be > 0") ++ for key in ( ++ "d_model", ++ "mlp_hidden", ++ "expert_hidden", ++ "shared_expert_hidden", ++ "delta_key_dim", ++ "delta_value_dim", ++ ): ++ cfg[key] = max(1, int(round(int(cfg[key]) * mult))) ++ attn_heads = int(cfg["attn_heads"]) ++ head_dim = int(DEFAULTS["d_model"]) // int(DEFAULTS["attn_heads"]) ++ if head_dim > 0 and cfg["d_model"] % head_dim == 0: ++ cfg["attn_heads"] = cfg["d_model"] // head_dim ++ elif cfg["d_model"] % attn_heads != 0: ++ h = min(attn_heads, cfg["d_model"]) ++ while h > 1 and cfg["d_model"] % h != 0: ++ h -= 1 ++ cfg["attn_heads"] = h ++ return cfg ++ ++ ++def build_loopmoe(ctx): ++ ctx = ctx if isinstance(ctx, dict) else {} ++ torch.manual_seed(int(ctx.get("seed", 0))) ++ te_flag = bool(ctx.get("te_available", False)) ++ if not te_flag: ++ te_flag = _probe_te_linear() is not None ++ return LoopMoE(_config_from_ctx(ctx), use_te=te_flag) +diff --git a/requirements.txt b/requirements.txt +new file mode 100644 +index 00000000..f5ccf09f +--- /dev/null ++++ b/requirements.txt +@@ -0,0 +1,12 @@ ++# CUDA 13 / cp312. Do not use [pytorch] extra: it pulls the PyPI ++# transformer_engine_torch==2.16.0 sdist (needs nvcc; Lium has none). ++# Do not pin the Astral local version: wheels.astral.sh connect-times ++# out from Lium (seen on 4061a769). Use NVIDIA's GitHub release wheel ++# (github.com / release-assets.githubusercontent.com) + PyPI core_cu13. ++# FLA: pin fla-core wheel + einops ONLY. Never flash-linear-attention or ++# fla-core[cuda] — those extras pull torch rebuilds / TE sdist. ++# Do not put pip option lines here. ++transformer-engine[core_cu13]==2.16.0 ++transformer_engine_torch @ https://github.com/NVIDIA/TransformerEngine/releases/download/v2.16/transformer_engine_torch-2.16.0+cu13torch26.05cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ++einops ++fla-core==0.5.2 diff --git a/examples/loopmoe/ddp_worker.py b/examples/loopmoe/ddp_worker.py new file mode 100644 index 0000000..41c9e89 --- /dev/null +++ b/examples/loopmoe/ddp_worker.py @@ -0,0 +1,47 @@ +"""LoopMoE DDP worker — one process per GPU via torch.multiprocessing.spawn. + +Rendezvous is tcp://127.0.0.1 (never hostname localhost — AF_INET6 errno 97). +No socket/subprocess imports (intake static_source NetworkExfil). + +submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z +""" + +from __future__ import annotations + +import os + +import torch.multiprocessing as mp + + +def _entry(rank, world, port, payload_path): + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world) + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + os.environ["LOOPMOE_PAYLOAD"] = payload_path + # Isolate Triton compile cache per rank (FLA autotune races under spawn). + os.environ["TRITON_CACHE_DIR"] = f"/tmp/loopmoe_triton_r{rank}" + os.environ.setdefault("LOOPMOE_PARALLEL", os.environ.get("LOOPMOE_PARALLEL", "ddp")) + os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo") + os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") + os.environ.setdefault("NCCL_IB_DISABLE", "1") + os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET") + from nemo_automodel.components.models.loopmoe.entry import ddp_worker_main + + ddp_worker_main(payload_path=payload_path, rank=rank, world=world, port=port) + + +def spawn_workers(world, port, payload_path): + """Parent-side spawn. Children re-import this module (real package path).""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo") + os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") + os.environ.setdefault("NCCL_IB_DISABLE", "1") + os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET") + try: + mp.set_start_method("spawn", force=True) + except RuntimeError: + pass + mp.spawn(_entry, nprocs=int(world), args=(int(world), int(port), str(payload_path)), join=True) diff --git a/examples/loopmoe/entry.py b/examples/loopmoe/entry.py new file mode 100644 index 0000000..b669b34 --- /dev/null +++ b/examples/loopmoe/entry.py @@ -0,0 +1,846 @@ +"""Prism-shaped AutoModel entry for LoopMoE (recipe 2.0). + +Exposes build_model / train for the operator harness seams. Uses FineWeb +stream + prism_telemetry from ctx. Real single-node DDP (one process per +GPU) + NVFP4 TE recipe when the class exists. + +submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z +""" + +from __future__ import annotations + +import json +import math +import os +import time +from datetime import timedelta +from pathlib import Path + +import torch + +from nemo_automodel.components.models.loopmoe import kernels as loopmoe_kernels +from nemo_automodel.components.models.loopmoe.model import build_loopmoe + +try: + import prism_telemetry +except ImportError: + + class _TelemetryFallback: + @staticmethod + def report(**_kwargs): + return None + + @staticmethod + def finish_evaluation(): + return None + + prism_telemetry = _TelemetryFallback() + + +PEAK_LR = 3e-4 +WEIGHT_DECAY = 0.1 +BETAS = (0.9, 0.95) +EPS = 1e-8 +WARMUP_FRAC = 0.02 +MIN_LR_FRAC = 0.10 +GRAD_CLIP = 1.0 +AUX_LOSS_COEF = 0.01 +REPORT_EVERY = 10 +WALL_MARGIN_S = 90.0 +# Per-GPU microbatch. Harness default is 8 *then DataParallel-sharded*. +# DDP keeps this whole batch on every rank (× world_size global tokens). +# Factored WY drops the 5-D decay tensor; mb=8 feeds GEMMs (seq stays 512). +DEFAULT_MICRO_BATCH = 8 +PEAK_FLOPS_PER_GPU = 209.5e12 +PAYLOAD_NAME = "loopmoe_ddp_payload.pt" +METRICS_NAME = "loopmoe_ddp_metrics.json" +WEIGHTS_NAME = "loopmoe_ddp_weights.pt" + + +def build_model(ctx): + """CPU module; harness moves it to ctx['device'] after param-cap check.""" + return build_loopmoe(ctx) + + +def _param_groups(model): + decay, no_decay = [], [] + for name, p in model.named_parameters(): + if not p.requires_grad: + continue + if p.ndim < 2 or "emb" in name or "loop_bias" in name or "inject_scale" in name: + no_decay.append(p) + else: + decay.append(p) + return [ + {"params": decay, "weight_decay": WEIGHT_DECAY}, + {"params": no_decay, "weight_decay": 0.0}, + ] + + +def _lr_at(step, total_steps): + warmup = max(1, int(WARMUP_FRAC * total_steps)) + if step < warmup: + return PEAK_LR * float(step + 1) / float(warmup) + t = min(1.0, (step - warmup) / max(1, total_steps - warmup)) + cos = 0.5 * (1.0 + math.cos(math.pi * t)) + return PEAK_LR * (MIN_LR_FRAC + (1.0 - MIN_LR_FRAC) * cos) + + +def _maybe_te_recipe(): + """Prefer NVFP4. SM120 (consumer Blackwell) disables RHT / stochastic rounding.""" + try: + from transformer_engine.common import recipe as te_recipe # type: ignore + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] TE recipe import failed ({exc})", flush=True) + return None, "none" + sm = None + if torch.cuda.is_available(): + sm = torch.cuda.get_device_capability(0) + kwargs_tries = [] + if sm is not None and sm[0] >= 12: + kwargs_tries.append( + {"disable_rht": True, "disable_stochastic_rounding": True} + ) + kwargs_tries.append({"disable_rht": True}) + kwargs_tries.append({}) + for name in ("NVFP4BlockScaling", "Float4BlockScaling", "MXFP4BlockScaling"): + cls = getattr(te_recipe, name, None) + if cls is None: + continue + for kw in kwargs_tries: + try: + rec = cls(**kw) + print( + f"[loopmoe] NVFP4 recipe class={name} kwargs={kw} sm={sm}", + flush=True, + ) + return rec, "nvfp4" + except TypeError: + continue + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] {name}({kw}) failed ({exc})", flush=True) + continue + delayed = getattr(te_recipe, "DelayedScaling", None) + if delayed is not None: + try: + print("[loopmoe] NVFP4 class missing; DelayedScaling FP8 fallback", flush=True) + return delayed(), "fp8" + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] DelayedScaling failed ({exc})", flush=True) + return None, "none" + + +def _fp8_ctx(enabled, rec): + """TE 2.16+ uses `autocast`; older wheels still export `fp8_autocast`.""" + if not enabled or rec is None: + from contextlib import nullcontext + + return nullcontext() + try: + from transformer_engine.pytorch import autocast as te_autocast # type: ignore + + try: + ctx = te_autocast(enabled=True, recipe=rec) + except TypeError: + ctx = te_autocast(recipe=rec) + if not getattr(_fp8_ctx, "_logged", False): + print("[loopmoe] using te.autocast for NVFP4/FP8 recipe (fwd+bwd)", flush=True) + _fp8_ctx._logged = True + return ctx + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] te.autocast unavailable ({exc}); trying fp8_autocast", flush=True) + try: + from transformer_engine.pytorch import fp8_autocast # type: ignore + + if not getattr(_fp8_ctx, "_logged", False): + print("[loopmoe] using te.fp8_autocast for NVFP4/FP8 recipe", flush=True) + _fp8_ctx._logged = True + return fp8_autocast(enabled=True, fp8_recipe=rec) + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] fp8_autocast unavailable ({exc}); BF16", flush=True) + from contextlib import nullcontext + + return nullcontext() + + +def _unwrap(model): + return model.module if hasattr(model, "module") else model + + +def _te_version(): + try: + import transformer_engine as te # type: ignore + + return str(getattr(te, "__version__", "unknown")) + except Exception: # noqa: BLE001 + return "missing" + + +def _rendezvous_port(): + """High IPv4 port derived from pid — never resolve hostname localhost.""" + return 29511 + (os.getpid() % 487) + + +def _set_dist_env(port): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + os.environ.setdefault("NCCL_SOCKET_IFNAME", "lo") + os.environ.setdefault("GLOO_SOCKET_IFNAME", "lo") + os.environ.setdefault("NCCL_IB_DISABLE", "1") + os.environ.setdefault("NCCL_SOCKET_FAMILY", "AF_INET") + os.environ.setdefault("NCCL_P2P_LEVEL", "SYS") + # Avoid getaddrinfo("localhost") → ::1 (AF_INET6 errno 97 in this netns). + os.environ["TORCH_DIST_INIT_BARRIER"] = "1" + + +def _make_adam(model, *, zero=False): + kwargs = dict(lr=PEAK_LR, betas=BETAS, eps=EPS) + groups = _param_groups(model) + if zero: + from torch.distributed.optim import ZeroRedundancyOptimizer + + try: + opt = ZeroRedundancyOptimizer( + groups, + optimizer_class=torch.optim.AdamW, + parameters_as_bucket_view=True, + fused=True, + **kwargs, + ) + print("[loopmoe] ZeRO-1 ZeroRedundancyOptimizer fused AdamW", flush=True) + return opt + except (TypeError, RuntimeError) as exc: + print(f"[loopmoe] ZeRO-1 fused failed ({exc}); plain AdamW", flush=True) + try: + return ZeroRedundancyOptimizer( + groups, + optimizer_class=torch.optim.AdamW, + parameters_as_bucket_view=True, + **kwargs, + ) + except Exception as exc2: # noqa: BLE001 + print(f"[loopmoe] ZeRO-1 unavailable ({exc2}); DDP AdamW", flush=True) + try: + return torch.optim.AdamW(groups, fused=True, **kwargs) + except (TypeError, RuntimeError) as exc: + print(f"[loopmoe] fused AdamW unavailable ({exc}); foreach", flush=True) + try: + return torch.optim.AdamW(groups, foreach=True, **kwargs) + except TypeError: + return torch.optim.AdamW(groups, **kwargs) + + +def _maybe_compile(model): + # TE NVFP4 + dynamo OOMed the first 8-GPU smoke; enable only when + # LOOPMOE_COMPILE=1 after a saturated eager run. + if os.environ.get("LOOPMOE_COMPILE", "").strip() not in {"1", "true", "yes"}: + return model, False + compile_fn = getattr(torch, "compile", None) + if compile_fn is None: + return model, False + try: + compiled = compile_fn(model, mode="default", fullgraph=False, dynamic=False) + return compiled, True + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] torch.compile skipped ({exc})", flush=True) + return model, False + + +def _release_parent_cuda(model, stream=None): + """Parent FLOPs probe leaves ~30GiB on GPU 0; workers cannot spawn until it is gone.""" + import gc + + def _cpu_tensors(obj): + for child in obj.modules() if hasattr(obj, "modules") else []: + for name, val in list(vars(child).items()): + if torch.is_tensor(val) and val.is_cuda: + setattr(child, name, val.detach().cpu()) + for p in obj.parameters(): + p.grad = None + if p.data.is_cuda: + p.data = p.data.cpu() + for b in obj.buffers(): + if b.is_cuda: + b.data = b.data.cpu() + + model.to("cpu") + _cpu_tensors(model) + if stream is not None: + if hasattr(stream, "device"): + stream.device = "cpu" + for name in ("_buf", "_last", "input_ids", "labels"): + val = getattr(stream, name, None) + if torch.is_tensor(val) and val.is_cuda: + setattr(stream, name, val.detach().cpu()) + gc.collect() + if torch.cuda.is_available(): + try: + torch.cuda.synchronize() + except Exception: # noqa: BLE001 + pass + torch.cuda.empty_cache() + try: + torch.cuda.ipc_collect() + except Exception: # noqa: BLE001 + pass + freed = [] + for i in range(torch.cuda.device_count()): + try: + free, total = torch.cuda.mem_get_info(i) + freed.append(f"{i}:{free/1e9:.2f}/{total/1e9:.2f}GiB") + except Exception: # noqa: BLE001 + continue + print(f"[loopmoe] parent CUDA released {freed}", flush=True) + + +class _LocalStream: + """Rank-local FineWeb stream (same contract as harness SeededTrainStream).""" + + def __init__(self, texts, tok, device, seq_len, batch_size, seed, rank=0): + self._texts = list(texts) + if not self._texts: + raise ValueError("empty train text pool") + self._tok = tok + self.device = device + self.seq_len = max(8, int(seq_len)) + self.batch_size = max(1, int(batch_size)) + self.seed = int(seed) + 10007 * int(rank) + self.tokens_seen = 0 + self._epoch = 0 + self._order = self._perm(0) + self._pos = 0 + self._buf = [] + self._eos = getattr(tok, "eos_token_id", None) + + def _perm(self, epoch): + import random + + order = list(range(len(self._texts))) + random.Random(self.seed + epoch).shuffle(order) + return order + + def _encode(self, text): + return self._tok(text, add_special_tokens=False)["input_ids"] + + def _fill(self): + need = self.batch_size * (self.seq_len + 1) + while len(self._buf) < need: + if self._pos >= len(self._order): + self._epoch += 1 + self._order = self._perm(self._epoch) + self._pos = 0 + text = self._texts[self._order[self._pos]] + self._pos += 1 + ids = self._encode(text) + if not ids: + continue + self._buf.extend(ids) + if self._eos is not None: + self._buf.append(self._eos) + + def next_batch(self): + self._fill() + need = self.batch_size * (self.seq_len + 1) + window = self._buf[:need] + del self._buf[:need] + ids = torch.tensor(window, dtype=torch.long).view(self.batch_size, self.seq_len + 1) + input_ids = ids[:, :-1].contiguous().to(self.device, non_blocking=True) + labels = ids[:, 1:].contiguous().to(self.device, non_blocking=True) + self.tokens_seen += int(labels.numel()) + return input_ids, labels + + +def _enable_fast_matmul(): + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + try: + torch.set_float32_matmul_precision("high") + except Exception: # noqa: BLE001 + pass + + +def _train_loop( + train_model, + stream, + *, + device, + max_steps, + stop_s, + rec, + te_mode, + guard, + rank=0, + world=1, + zero=False, +): + """Single backward: CE + local MoE aux. No second backward, no DP gather.""" + core = _unwrap(train_model) + if hasattr(core, "grad_checkpoint"): + # TE NVFP4 Linear cannot recompute under torch.utils.checkpoint + # (saved-tensor count 94 vs 45). Aux is already in the same loss. + core.grad_checkpoint = False + opt = _make_adam(core, zero=zero) + use_amp = device == "cuda" + use_te = rec is not None + t0 = time.time() + step = 0 + last_loss = 0.0 + last_aux = 0.0 + grad_norm = 0.0 + tokens_this = 0 + train_model.train() + while step < max_steps and (time.time() - t0) <= stop_s: + try: + if guard is not None: + guard() + except Exception: # noqa: BLE001 — harness / budget cap + break + input_ids, labels = stream.next_batch() if hasattr(stream, "next_batch") else next(stream) + tokens_this += int(input_ids.numel()) + # TE recipe MUST wrap backward — closing autocast after forward + # makes NVFP4 wgrad pick a cublasLt algo that SM120 rejects. + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_amp): + with _fp8_ctx(use_te, rec): + logits = train_model(input_ids) + if hasattr(logits, "logits"): + logits = logits.logits + loss = loopmoe_kernels.cross_entropy( + logits.float().reshape(-1, logits.shape[-1]), labels.reshape(-1) + ) + aux = getattr(core, "aux_loss", None) + if aux is not None and torch.is_tensor(aux) and aux.requires_grad: + last_aux = float(aux.detach().float().item()) + loss = loss + AUX_LOSS_COEF * aux.float() + elif aux is not None and torch.is_tensor(aux): + last_aux = float(aux.detach().float().item()) + opt.zero_grad(set_to_none=True) + loss.backward() + grad_norm = float(torch.nn.utils.clip_grad_norm_(core.parameters(), GRAD_CLIP)) + lr = _lr_at(step, max_steps) + for group in opt.param_groups: + group["lr"] = lr + opt.step() + last_loss = float(loss.detach().float().item()) + step += 1 + if rank == 0 and step == 1 and torch.cuda.is_available(): + try: + free, total = torch.cuda.mem_get_info() + km = loopmoe_kernels.kernel_map() + print( + f"[loopmoe] step1 mem_free={free/1e9:.2f}/{total/1e9:.2f}GiB " + f"ckpt={getattr(core, 'grad_checkpoint', None)} te_mode={te_mode} " + f"delta_kernel={km.get('delta_kernel')} attn_kernel={km.get('attn_kernel')} " + f"ce_kernel={km.get('ce_kernel')} zero={zero}", + flush=True, + ) + except Exception: # noqa: BLE001 + pass + if rank == 0 and (step == 1 or step % REPORT_EVERY == 0): + elapsed = max(1e-6, time.time() - t0) + tps_local = tokens_this / elapsed + tps_global = tps_local * world + km = loopmoe_kernels.kernel_map() + print( + f"[loopmoe] step={step} loss={last_loss:.4f} aux={last_aux:.4f} " + f"tok/s_local={tps_local:.1f} tok/s_global={tps_global:.1f} " + f"world={world} te_mode={te_mode} rank={rank} " + f"delta_kernel={km.get('delta_kernel')} attn_kernel={km.get('attn_kernel')}", + flush=True, + ) + prism_telemetry.report(loss=last_loss, step=step, grad_norm=grad_norm) + elapsed = time.time() - t0 + tps_local = tokens_this / max(1e-6, elapsed) + return { + "train_loss": last_loss, + "train_steps": step, + "train_seconds": elapsed, + "moe_aux_loss": last_aux, + "tokens_local": tokens_this, + "tokens_per_sec_local": tps_local, + "tokens_per_sec": tps_local * world, + "final_lr": _lr_at(max(step - 1, 0), max_steps), + "peak_lr": PEAK_LR, + } + + +def ddp_worker_main(payload_path=None, rank=None, world=None, port=None): + """One process per GPU. Called from ddp_worker.py via mp.spawn.""" + payload_path = payload_path or os.environ.get("LOOPMOE_PAYLOAD") + rank = int(os.environ["RANK"] if rank is None else rank) + world = int(os.environ["WORLD_SIZE"] if world is None else world) + port = int(os.environ["MASTER_PORT"] if port is None else port) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + _set_dist_env(port) + torch.cuda.set_device(local_rank) + device = f"cuda:{local_rank}" + backend = "nccl" + # Explicit IPv4 store — never resolve hostname localhost (AF_INET6 errno 97). + store_kwargs = dict( + host_name="127.0.0.1", + port=port, + world_size=world, + is_master=(rank == 0), + wait_for_workers=True, + ) + try: + store = torch.distributed.TCPStore(**store_kwargs, use_libuv=False) + except TypeError: + store = torch.distributed.TCPStore("127.0.0.1", port, world, rank == 0, True) + torch.distributed.init_process_group( + backend=backend, + store=store, + rank=rank, + world_size=world, + timeout=timedelta(minutes=15), + ) + print( + f"[loopmoe] ddp init rank={rank}/{world} local_rank={local_rank} " + f"backend={backend} master=127.0.0.1:{port} " + f"device={torch.cuda.get_device_name(local_rank)} " + f"sm={torch.cuda.get_device_capability(local_rank)} " + f"nccl={getattr(torch.cuda.nccl, 'version', lambda: '?')()}", + flush=True, + ) + print(f"[loopmoe] rank={rank} loading payload", flush=True) + payload = torch.load(payload_path, map_location="cpu", weights_only=False) + ctx = dict(payload["ctx"]) + ctx["device"] = device + ctx["te_available"] = True + texts = list(payload.get("texts") or []) + texts_path = payload.get("texts_path") + if not texts and texts_path: + texts = [json.loads(line) for line in open(texts_path, encoding="utf-8") if line.strip()] + if not texts: + raise RuntimeError("DDP worker missing train texts") + print(f"[loopmoe] rank={rank} texts={len(texts)} building model", flush=True) + model = build_loopmoe(ctx) + # TE Linear writes `_extra_state` during the parent FLOPs probe; a + # freshly constructed worker module does not declare those keys yet. + model.load_state_dict(payload["state_dict"], strict=False) + model = model.to(device) + micro = int(payload["micro_batch"]) + seq_len = int(payload["seq_len"]) + rec, te_mode = _maybe_te_recipe() + parallel = str(payload.get("parallel") or os.environ.get("LOOPMOE_PARALLEL", "ddp")).strip().lower() + if parallel not in {"ddp", "zero1", "fsdp"}: + parallel = "ddp" + print( + f"[loopmoe] worker te_version={_te_version()} te_mode={te_mode} " + f"use_te_linear={getattr(model, 'use_te', None)} parallel={parallel}", + flush=True, + ) + loopmoe_kernels.enable_attn_backends() + loopmoe_kernels.log_kernel_banner() + use_zero = parallel == "zero1" + if parallel == "fsdp": + train_wrap, parallel_used = _wrap_fsdp(model, local_rank) + else: + train_wrap = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[local_rank], + output_device=local_rank, + broadcast_buffers=True, + find_unused_parameters=False, + gradient_as_bucket_view=True, + ) + parallel_used = "zero1" if use_zero else "ddp" + compiled, did_compile = _maybe_compile(train_wrap) + _enable_fast_matmul() + stream = _LocalStream( + texts, + payload["tokenizer"], + device, + seq_len=seq_len, + batch_size=micro, + seed=int(payload["seed"]), + rank=rank, + ) + cap_s = float(payload["cap_s"]) + stop_s = max(30.0, cap_s - float(payload.get("wall_margin_s", WALL_MARGIN_S))) + t_limit = time.time() + stop_s + + def guard(): + if time.time() >= t_limit: + raise RuntimeError("wall") + + metrics = _train_loop( + compiled, + stream, + device="cuda", + max_steps=int(payload["max_steps"]), + stop_s=stop_s, + rec=rec, + te_mode=te_mode, + guard=guard, + rank=rank, + world=world, + zero=use_zero, + ) + km = loopmoe_kernels.kernel_map() + metrics.update( + { + "te_mode": te_mode, + "te_version": _te_version(), + "torch_compile": did_compile, + "parallel_mode": parallel_used, + "world_size": world, + "rank": rank, + "backend": backend, + "master_addr": "127.0.0.1", + "micro_batch": micro, + "seq_len": seq_len, + "gpu_count": world, + "te_available": True, + **km, + } + ) + tokens_t = torch.tensor([float(metrics["tokens_local"])], device=device) + torch.distributed.all_reduce(tokens_t, op=torch.distributed.ReduceOp.SUM) + metrics["tokens_seen"] = int(tokens_t.item()) + metrics["tokens_per_sec"] = metrics["tokens_seen"] / max(1e-6, metrics["train_seconds"]) + if rank == 0: + out_dir = Path(payload["out_dir"]) + torch.save({k: v.detach().cpu() for k, v in _unwrap(compiled).state_dict().items()}, out_dir / WEIGHTS_NAME) + (out_dir / METRICS_NAME).write_text(json.dumps(metrics, indent=2), encoding="utf-8") + print( + f"[loopmoe] train done steps={metrics['train_steps']} " + f"seconds={metrics['train_seconds']:.1f} loss={metrics['train_loss']:.4f} " + f"tokens={metrics['tokens_seen']} tok/s={metrics['tokens_per_sec']:.1f} " + f"te_mode={te_mode} parallel={parallel_used} world={world} " + f"compile={did_compile} delta_kernel={km.get('delta_kernel')}", + flush=True, + ) + torch.distributed.barrier() + torch.distributed.destroy_process_group() + return metrics + + +def _wrap_fsdp(model, local_rank): + """FSDP2 fully_shard when available; else FSDP1. TE Linear extra-state is sticky.""" + try: + from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard + + mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) + for child in list(model.children()): + try: + fully_shard(child, mp_policy=mp) + except Exception: # noqa: BLE001 + continue + fully_shard(model, mp_policy=mp) + print(f"[loopmoe] FSDP2 fully_shard rank={local_rank}", flush=True) + return model, "fsdp2" + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] FSDP2 unavailable ({exc}); trying FSDP1", flush=True) + try: + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp import MixedPrecision + + mp = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.bfloat16, + ) + wrapped = FSDP( + model, + mixed_precision=mp, + use_orig_params=True, + device_id=local_rank, + ) + print(f"[loopmoe] FSDP1 wrap rank={local_rank}", flush=True) + return wrapped, "fsdp1" + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] FSDP failed ({exc}); falling back to DDP", flush=True) + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[local_rank], + output_device=local_rank, + broadcast_buffers=True, + find_unused_parameters=False, + gradient_as_bucket_view=True, + ) + return ddp, "ddp" + + +def _launch_ddp(model, ctx, gpu_count): + workdir = Path(ctx.get("workdir") or os.environ.get("PRISM_WORKDIR") or "/tmp") + out_dir = workdir / "loopmoe_ddp" + out_dir.mkdir(parents=True, exist_ok=True) + stream = ctx.get("train_stream") + if stream is None: + raise RuntimeError("train_stream required for DDP LoopMoE") + texts = list(getattr(stream, "_texts", []) or []) + tok = ctx.get("tokenizer") or getattr(stream, "_tok", None) + if not texts or tok is None: + raise RuntimeError("DDP payload needs stream texts + tokenizer") + # Free parent CUDA so workers own the devices (probe left ~30GiB on GPU 0). + cpu_sd = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()} + _release_parent_cuda(model, stream) + seq_len = int(ctx.get("seq_len") or getattr(stream, "seq_len", 512) or 512) + harness_bs = int(ctx.get("batch_size") or getattr(stream, "batch_size", 8) or 8) + env_micro = os.environ.get("LOOPMOE_MICRO_BATCH", "").strip() + # Do not inherit harness batch_size (that was DP-sharded). LoopMoE + # activations at seq=512 need a small per-GPU microbatch. + micro = int(env_micro) if env_micro.isdigit() else DEFAULT_MICRO_BATCH + _ = harness_bs # kept for payload logs / MFU context + cap_s = float(ctx.get("train_hours_cap", 1.0)) * 3600.0 + texts_path = out_dir / "train_texts.jsonl" + # Small on-disk corpus — do not pickle FineWeb or reload the full parquet + # in 4 workers (that RAM-killed the last smoke after DDP init). + with open(texts_path, "w", encoding="utf-8") as fh: + for text in texts[:4096]: + fh.write(json.dumps(text, ensure_ascii=False) + "\n") + payload = { + "state_dict": cpu_sd, + "texts": [], + "texts_path": str(texts_path), + "tokenizer": tok, + "ctx": { + "seed": int(ctx.get("seed", 0)), + "vocab_size": int(ctx.get("vocab_size") or 50257), + "te_available": True, + "arch": ctx.get("arch"), + "prism_width_multiplier": ctx.get("prism_width_multiplier", 1.0), + }, + "seed": int(ctx.get("seed", 0)), + "seq_len": seq_len, + "micro_batch": micro, + "max_steps": int(ctx.get("max_train_steps", 20000)), + "cap_s": cap_s, + "wall_margin_s": WALL_MARGIN_S, + "out_dir": str(out_dir), + "parallel": os.environ.get("LOOPMOE_PARALLEL", "ddp").strip().lower(), + } + payload_path = out_dir / PAYLOAD_NAME + torch.save(payload, payload_path) + port = _rendezvous_port() + _set_dist_env(port) + print( + f"[loopmoe] launching dist spawn world={gpu_count} master=127.0.0.1:{port} " + f"micro_batch={micro} seq={seq_len} parallel={payload['parallel']}", + flush=True, + ) + from nemo_automodel.components.models.loopmoe.ddp_worker import spawn_workers + + spawn_workers(gpu_count, port, str(payload_path)) + metrics_path = out_dir / METRICS_NAME + weights_path = out_dir / WEIGHTS_NAME + if not metrics_path.is_file() or not weights_path.is_file(): + raise RuntimeError("DDP workers did not write metrics/weights") + metrics = json.loads(metrics_path.read_text(encoding="utf-8")) + trained = torch.load(weights_path, map_location="cpu", weights_only=False) + model.load_state_dict(trained, strict=False) + device = ctx.get("device") or "cuda" + if device != "cpu": + model.to(device) + # Authoritative harness counter + FLOPs spend so MFU is real. + tokens = int(metrics.get("tokens_seen") or 0) + if hasattr(stream, "tokens_seen"): + stream.tokens_seen = int(getattr(stream, "tokens_seen", 0)) + tokens + if getattr(stream, "flops_per_token", 0.0): + stream.flops_spent = float(stream.flops_per_token) * float(stream.tokens_seen) + stream.batches_yielded = int(getattr(stream, "batches_yielded", 0)) + int( + metrics.get("train_steps") or 0 + ) + return metrics + + +def train(model, ctx): + device = ctx.get("device") or "cuda" + seed = int(ctx.get("seed", 0)) + torch.manual_seed(seed) + if device == "cuda" and torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + _enable_fast_matmul() + guard = ctx.get("guard") + gpu_count = int(ctx.get("gpu_count") or 0) + if gpu_count <= 0 and torch.cuda.is_available(): + gpu_count = int(torch.cuda.device_count()) + te_available = bool(ctx.get("te_available", False)) + if not te_available: + try: + te_available = __import__("importlib").util.find_spec("transformer_engine") is not None + except Exception: # noqa: BLE001 + te_available = False + + rec, te_mode = _maybe_te_recipe() if te_available else (None, "none") + print( + f"[loopmoe] train start gpu_count={gpu_count} te_available={te_available} " + f"te_mode={te_mode} te_version={_te_version()} " + f"cuda_devices={torch.cuda.device_count() if torch.cuda.is_available() else 0} " + f"use_te_linear={getattr(model, 'use_te', None)}", + flush=True, + ) + + # Marketplace often only lists 8×5090 hosts (no GPU splitting). Cap at 4 + # so the proof matches the 4-GPU contract and leaves headroom on GPU 0. + max_gpus = int(os.environ.get("LOOPMOE_MAX_GPUS", "4") or 4) + if gpu_count > max_gpus: + print(f"[loopmoe] capping visible GPUs {gpu_count} -> {max_gpus}", flush=True) + gpu_count = max_gpus + if gpu_count > 1 and torch.cuda.is_available(): + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(gpu_count)) + metrics = _launch_ddp(model, ctx, gpu_count) + fpt = float(ctx.get("flops_per_token_probe") or 0.0) + elapsed = float(metrics.get("train_seconds") or 1.0) + tokens = float(metrics.get("tokens_seen") or 0.0) + n_params = float(sum(p.numel() for p in model.parameters())) + loop_f = float(getattr(model, "prism_loop_factor", 1.0) or 1.0) + fpt_analytic = 6.0 * n_params * loop_f + fpt_source = "probe" + if fpt <= 0.0: + fpt = fpt_analytic + fpt_source = "analytic_6n_loops" + mfu = (tokens * fpt) / (PEAK_FLOPS_PER_GPU * gpu_count * elapsed) if fpt > 0 else 0.0 + metrics["mfu_est"] = mfu + metrics["flops_per_token_probe"] = fpt + metrics["flops_per_token_analytic"] = fpt_analytic + metrics["flops_per_token_source"] = fpt_source + print( + f"[loopmoe] ddp parent metrics world={metrics.get('world_size')} " + f"te_mode={metrics.get('te_mode')} tok/s={metrics.get('tokens_per_sec')} " + f"mfu_est={mfu*100:.2f}% fpt_src={fpt_source} compile={metrics.get('torch_compile')}", + flush=True, + ) + prism_telemetry.finish_evaluation() + return metrics + + stream = ctx.get("train_stream") + if stream is None: + raise RuntimeError("train_stream required for live AutoModel LoopMoE") + compiled, did_compile = _maybe_compile(model) + max_steps = int(ctx.get("max_train_steps", 20000)) + cap_s = float(ctx.get("train_hours_cap", 1.0)) * 3600.0 + stop_s = max(60.0, cap_s - WALL_MARGIN_S) + metrics = _train_loop( + compiled, + stream, + device=device, + max_steps=max_steps, + stop_s=stop_s, + rec=rec, + te_mode=te_mode, + guard=guard, + rank=0, + world=1, + ) + metrics.update( + { + "te_mode": te_mode, + "te_version": _te_version(), + "te_available": te_available, + "torch_compile": did_compile, + "parallel_mode": "single", + "world_size": 1, + "backend": "none", + "gpu_count": gpu_count, + "tokens_seen": int(getattr(stream, "tokens_seen", 0)) or int(metrics["tokens_local"]), + } + ) + print( + f"[loopmoe] train done steps={metrics['train_steps']} " + f"seconds={metrics['train_seconds']:.1f} te_mode={te_mode} parallel=single", + flush=True, + ) + prism_telemetry.finish_evaluation() + return metrics diff --git a/examples/loopmoe/kernels.py b/examples/loopmoe/kernels.py new file mode 100644 index 0000000..61e6ea2 --- /dev/null +++ b/examples/loopmoe/kernels.py @@ -0,0 +1,867 @@ +"""Fused kernels for LoopMoE — gated-delta scan, attention, RMS, CE. + +submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z + +Default delta path is **in-pack factored chunked WY** (batched GEMMs). +Never the sequential-T Triton scan. Never auto-select FLA `chunk_kda` — +its Triton autotune cache races under DDP spawn (seen on 01f64b4f). + +Optional: + LOOPMOE_DELTA_KERNEL=kda — fla-core chunk_kda (per-rank Triton cache) + LOOPMOE_DELTA_KERNEL=triton — recurrent scan (slower; do not ship) + LOOPMOE_DELTA_KERNEL=eager — legacy per-chunk WY loop + +Hot-path train (seq=512 <= window) never walks `for t in range(seq)`. +""" + +from __future__ import annotations + +import math +import os + +import torch +import torch.nn.functional as F + +DELTA_KERNEL = "uninitialized" +ATTN_KERNEL = "sdpa" +RMS_KERNEL = "torch" +CE_KERNEL = "torch" +SWIGLU_KERNEL = "eager" +ROPE_KERNEL = "torch" + +_DELTA_LOGGED = False +_TRITON_OK = None +_KDA_FN = None +_KDA_PROBED = False +_GDR_FN = None +_GDR_PROBED = False +_KDA_CHECKED = False +_KDA_OK = False + + +def kernel_map(): + return { + "delta_kernel": DELTA_KERNEL, + "attn_kernel": ATTN_KERNEL, + "rmsnorm_kernel": RMS_KERNEL, + "ce_kernel": CE_KERNEL, + "swiglu_kernel": SWIGLU_KERNEL, + "rope_kernel": ROPE_KERNEL, + } + + +def _env_force(): + return os.environ.get("LOOPMOE_DELTA_KERNEL", "").strip().lower() + + +def _probe_kda(): + """FLA `chunk_kda` — per-channel decay, same recurrence as eager WY.""" + global _KDA_FN, _KDA_PROBED + if _KDA_PROBED: + return _KDA_FN + _KDA_PROBED = True + if _env_force() != "kda": + return None + try: + from fla.ops.kda import chunk_kda # type: ignore + + _KDA_FN = chunk_kda + except Exception: # noqa: BLE001 + _KDA_FN = None + return _KDA_FN + + +def _probe_gdr(): + """FLA `chunk_gated_delta_rule` — per-head scalar g only. Never default.""" + global _GDR_FN, _GDR_PROBED + if _GDR_PROBED: + return _GDR_FN + _GDR_PROBED = True + if _env_force() not in {"gdr", "fla"}: + return None + try: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule # type: ignore + + _GDR_FN = chunk_gated_delta_rule + except Exception: # noqa: BLE001 + _GDR_FN = None + return _GDR_FN + + +def _probe_triton(): + global _TRITON_OK + if _TRITON_OK is not None: + return _TRITON_OK + # Recurrent scan is opt-in only — do not even compile it on the default path. + if _env_force() != "triton": + _TRITON_OK = False + return False + try: + import triton # noqa: F401 + import triton.language as tl # noqa: F401 + + _TRITON_OK = bool(torch.cuda.is_available()) + except Exception: # noqa: BLE001 + _TRITON_OK = False + return _TRITON_OK + + +def _log_delta(name): + global DELTA_KERNEL, _DELTA_LOGGED + DELTA_KERNEL = name + if not _DELTA_LOGGED: + print(f"[loopmoe] delta_kernel={name}", flush=True) + _DELTA_LOGGED = True + + +# --------------------------------------------------------------------------- +# Eager sequential (reference) + batched WY (no token loop) +# --------------------------------------------------------------------------- + + +def _delta_sequential(q, k, v, beta, la): + """Exact recurrent gated-delta. q/k/v/beta/la float, heads flattened.""" + bh, t, dk = q.shape + dv = v.shape[-1] + state = q.new_zeros(bh, dv, dk) + outs = [] + for i in range(t): + alpha = la[:, i, :].exp() + state = state * alpha.unsqueeze(1) + kt = k[:, i, :] + vt = v[:, i, :] + bt = beta[:, i, :] + qt = q[:, i, :] + sk = torch.einsum("bvd,bd->bv", state, kt) + u = bt * (vt - sk) + state = state + u.unsqueeze(-1) * kt.unsqueeze(1) + outs.append(torch.einsum("bvd,bd->bv", state, qt)) + return torch.stack(outs, dim=1) + + +def _delta_chunk_loop(q, k, v, beta, la, chunk): + """Original WY chunk loop — last-resort fallback.""" + _, t, dk = q.shape + dv = v.shape[-1] + state = q.new_zeros(q.shape[0], dv, dk) + outs = [] + eye_full = torch.eye(chunk, device=q.device, dtype=q.dtype) + tril_full = torch.ones(chunk, chunk, dtype=torch.bool, device=q.device).tril() + for s in range(0, t, chunk): + e = min(s + chunk, t) + c = e - s + qc, kc, vc = q[:, s:e], k[:, s:e], v[:, s:e] + bc = beta[:, s:e] + L = la[:, s:e].cumsum(dim=1) + ldiff = L[:, :, None, :] - L[:, None, :, :] + dec = ldiff.masked_fill(~tril_full[:c, :c][None, :, :, None], float("-inf")).exp() + a = (bc * torch.einsum("btc,bic,btic->bti", kc, kc, dec)).tril(-1) + bm = torch.einsum("btc,bic,btic->bti", qc, kc, dec) + lexp = L.exp() + rhs = bc * (vc - (kc * lexp) @ state.transpose(-1, -2)) + u = torch.linalg.solve_triangular( + a + eye_full[:c, :c], rhs, upper=False, unitriangular=True + ) + outs.append((qc * lexp) @ state.transpose(-1, -2) + bm @ u) + e_lc = L[:, -1:, :].exp() + k_tail = kc * (L[:, -1:, :] - L).exp() + state = state * e_lc + u.transpose(-1, -2) @ k_tail + return torch.cat(outs, dim=1) + + +def _delta_vectorized(q, k, v, beta, la, chunk): + """Chunked WY via factored batched GEMMs (FLA algorithm, per-channel decay). + + Dec[t,i,d] = exp(L[t,d] − L[i,d]) is never materialized as a 5-D tensor. + Intra-chunk A/B are ``(x ⊙ e^L) @ (k ⊙ e^{−L})^T`` (cuBLAS), then a short + state-carry over n_chunks (16 at seq 512). Safe in fp32 for chunk<=40 + with |la|<=2 (exp(80) still finite). Longer chunks fall back to the + masked 5-D path. + """ + bh, t, dk = q.shape + dv = v.shape[-1] + n_chunks = (t + chunk - 1) // chunk + pad = n_chunks * chunk - t + if pad: + q = F.pad(q, (0, 0, 0, pad)) + k = F.pad(k, (0, 0, 0, pad)) + v = F.pad(v, (0, 0, 0, pad)) + beta = F.pad(beta, (0, 0, 0, pad)) + la = F.pad(la, (0, 0, 0, pad)) + qc = q.view(bh, n_chunks, chunk, dk) + kc = k.view(bh, n_chunks, chunk, dk) + vc = v.view(bh, n_chunks, chunk, dv) + bc = beta.view(bh, n_chunks, chunk, 1) + lac = la.view(bh, n_chunks, chunk, dk) + l = lac.cumsum(dim=2) + lexp = l.exp() + tril = torch.ones(chunk, chunk, dtype=torch.bool, device=q.device).tril() + if int(chunk) <= 40: + inv = (-l).exp() + k_neg_t = (kc * inv).transpose(-1, -2) + gram = torch.matmul(kc * lexp, k_neg_t) + a = (bc * gram).tril(-1) + bm = torch.matmul(qc * lexp, k_neg_t).masked_fill(~tril.view(1, 1, chunk, chunk), 0) + else: + ldiff = l.unsqueeze(3) - l.unsqueeze(2) + dec = ldiff.masked_fill(~tril.view(1, 1, chunk, chunk, 1), float("-inf")).exp() + a = (bc * torch.einsum("bntd,bnid,bntid->bnti", kc, kc, dec)).tril(-1) + bm = torch.einsum("bntd,bnid,bntid->bnti", qc, kc, dec) + eye = torch.eye(chunk, device=q.device, dtype=q.dtype) + state = q.new_zeros(bh, dv, dk) + outs = q.new_zeros(bh, n_chunks, chunk, dv) + # State carry across chunks — not a token loop. + for i in range(n_chunks): + rhs = bc[:, i] * (vc[:, i] - (kc[:, i] * lexp[:, i]) @ state.transpose(-1, -2)) + u = torch.linalg.solve_triangular(a[:, i] + eye, rhs, upper=False, unitriangular=True) + outs[:, i] = (qc[:, i] * lexp[:, i]) @ state.transpose(-1, -2) + bm[:, i] @ u + e_lc = lexp[:, i, -1:, :] + k_tail = kc[:, i] * (l[:, i, -1:, :] - l[:, i]).exp() + state = state * e_lc + u.transpose(-1, -2) @ k_tail + return outs.reshape(bh, n_chunks * chunk, dv)[:, :t] + + +# --------------------------------------------------------------------------- +# Triton fused recurrent (per-channel decay) + custom bwd +# --------------------------------------------------------------------------- + + +def _triton_kernels(): + import triton + import triton.language as tl + + @triton.jit + def fwd_kernel( + q_ptr, + k_ptr, + v_ptr, + b_ptr, + g_ptr, + o_ptr, + chk_ptr, + T, + DK, + DV, + stride_q_bh, + stride_q_t, + stride_q_d, + stride_k_bh, + stride_k_t, + stride_k_d, + stride_v_bh, + stride_v_t, + stride_v_d, + stride_b_bh, + stride_b_t, + stride_g_bh, + stride_g_t, + stride_g_d, + stride_o_bh, + stride_o_t, + stride_o_d, + stride_c_bh, + stride_c_n, + stride_c_v, + stride_c_k, + CHUNK: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_V: tl.constexpr, + ): + bh = tl.program_id(0) + offs_k = tl.arange(0, BLOCK_K) + offs_v = tl.arange(0, BLOCK_V) + mask_k = offs_k < DK + mask_v = offs_v < DV + mask_h = mask_v[:, None] & mask_k[None, :] + s = tl.zeros((BLOCK_V, BLOCK_K), dtype=tl.float32) + q_bh = q_ptr + bh * stride_q_bh + k_bh = k_ptr + bh * stride_k_bh + v_bh = v_ptr + bh * stride_v_bh + b_bh = b_ptr + bh * stride_b_bh + g_bh = g_ptr + bh * stride_g_bh + o_bh = o_ptr + bh * stride_o_bh + c_bh = chk_ptr + bh * stride_c_bh + tl.store( + c_bh + offs_v[:, None] * stride_c_v + offs_k[None, :] * stride_c_k, + s, + mask=mask_h, + ) + nchk = 1 + for t in range(0, T): + q = tl.load(q_bh + t * stride_q_t + offs_k * stride_q_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + k = tl.load(k_bh + t * stride_k_t + offs_k * stride_k_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + v = tl.load(v_bh + t * stride_v_t + offs_v * stride_v_d, mask=mask_v, other=0.0).to( + tl.float32 + ) + beta = tl.load(b_bh + t * stride_b_t).to(tl.float32) + gk = tl.load(g_bh + t * stride_g_t + offs_k * stride_g_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + alpha = tl.exp(gk) + s = s * alpha[None, :] + sk = tl.sum(s * k[None, :], axis=1) + u = beta * (v - sk) + s = s + u[:, None] * k[None, :] + o = tl.sum(s * q[None, :], axis=1) + tl.store( + o_bh + t * stride_o_t + offs_v * stride_o_d, o.to(o_ptr.dtype.element_ty), mask=mask_v + ) + if (t + 1) % CHUNK == 0: + tl.store( + c_bh + + nchk * stride_c_n + + offs_v[:, None] * stride_c_v + + offs_k[None, :] * stride_c_k, + s, + mask=mask_h, + ) + nchk = nchk + 1 + + @triton.jit + def bwd_kernel( + q_ptr, + k_ptr, + v_ptr, + b_ptr, + g_ptr, + do_ptr, + chk_ptr, + scratch_ptr, + dq_ptr, + dk_ptr, + dv_ptr, + db_ptr, + dg_ptr, + T, + NCHK, + DK, + DV, + stride_q_bh, + stride_q_t, + stride_q_d, + stride_k_bh, + stride_k_t, + stride_k_d, + stride_v_bh, + stride_v_t, + stride_v_d, + stride_b_bh, + stride_b_t, + stride_g_bh, + stride_g_t, + stride_g_d, + stride_o_bh, + stride_o_t, + stride_o_d, + stride_c_bh, + stride_c_n, + stride_c_v, + stride_c_k, + stride_s_bh, + stride_s_j, + stride_s_v, + stride_s_k, + CHUNK: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_V: tl.constexpr, + ): + bh = tl.program_id(0) + offs_k = tl.arange(0, BLOCK_K) + offs_v = tl.arange(0, BLOCK_V) + mask_k = offs_k < DK + mask_v = offs_v < DV + mask_h = mask_v[:, None] & mask_k[None, :] + q_bh = q_ptr + bh * stride_q_bh + k_bh = k_ptr + bh * stride_k_bh + v_bh = v_ptr + bh * stride_v_bh + b_bh = b_ptr + bh * stride_b_bh + g_bh = g_ptr + bh * stride_g_bh + do_bh = do_ptr + bh * stride_o_bh + c_bh = chk_ptr + bh * stride_c_bh + sc_bh = scratch_ptr + bh * stride_s_bh + dq_bh = dq_ptr + bh * stride_q_bh + dk_bh = dk_ptr + bh * stride_k_bh + dv_bh = dv_ptr + bh * stride_v_bh + db_bh = db_ptr + bh * stride_b_bh + dg_bh = dg_ptr + bh * stride_g_bh + ds = tl.zeros((BLOCK_V, BLOCK_K), dtype=tl.float32) + for ic in range(0, NCHK): + ci = NCHK - 1 - ic + s = tl.load( + c_bh + + ci * stride_c_n + + offs_v[:, None] * stride_c_v + + offs_k[None, :] * stride_c_k, + mask=mask_h, + other=0.0, + ).to(tl.float32) + for j in range(0, CHUNK): + t = ci * CHUNK + j + k = tl.load(k_bh + t * stride_k_t + offs_k * stride_k_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + v = tl.load(v_bh + t * stride_v_t + offs_v * stride_v_d, mask=mask_v, other=0.0).to( + tl.float32 + ) + beta = tl.load(b_bh + t * stride_b_t).to(tl.float32) + gk = tl.load(g_bh + t * stride_g_t + offs_k * stride_g_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + alpha = tl.exp(gk) + s = s * alpha[None, :] + sk = tl.sum(s * k[None, :], axis=1) + u = beta * (v - sk) + s = s + u[:, None] * k[None, :] + tl.store( + sc_bh + + j * stride_s_j + + offs_v[:, None] * stride_s_v + + offs_k[None, :] * stride_s_k, + s, + mask=mask_h, + ) + for jj in range(0, CHUNK): + j = CHUNK - 1 - jj + t = ci * CHUNK + j + q = tl.load(q_bh + t * stride_q_t + offs_k * stride_q_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + k = tl.load(k_bh + t * stride_k_t + offs_k * stride_k_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + v = tl.load(v_bh + t * stride_v_t + offs_v * stride_v_d, mask=mask_v, other=0.0).to( + tl.float32 + ) + beta = tl.load(b_bh + t * stride_b_t).to(tl.float32) + gk = tl.load(g_bh + t * stride_g_t + offs_k * stride_g_d, mask=mask_k, other=0.0).to( + tl.float32 + ) + dout = tl.load( + do_bh + t * stride_o_t + offs_v * stride_o_d, mask=mask_v, other=0.0 + ).to(tl.float32) + alpha = tl.exp(gk) + st = tl.load( + sc_bh + + j * stride_s_j + + offs_v[:, None] * stride_s_v + + offs_k[None, :] * stride_s_k, + mask=mask_h, + other=0.0, + ).to(tl.float32) + if j == 0: + s_prev = tl.load( + c_bh + + ci * stride_c_n + + offs_v[:, None] * stride_c_v + + offs_k[None, :] * stride_c_k, + mask=mask_h, + other=0.0, + ).to(tl.float32) + else: + s_prev = tl.load( + sc_bh + + (j - 1) * stride_s_j + + offs_v[:, None] * stride_s_v + + offs_k[None, :] * stride_s_k, + mask=mask_h, + other=0.0, + ).to(tl.float32) + s_mid = s_prev * alpha[None, :] + smk = tl.sum(s_mid * k[None, :], axis=1) + r = v - smk + u = beta * r + dq = tl.sum(st * dout[:, None], axis=0) + ds = ds + dout[:, None] * q[None, :] + du = tl.sum(ds * k[None, :], axis=1) + dk = tl.sum(ds * u[:, None], axis=0) + dbeta = tl.sum(du * r) + dr = du * beta + dk = dk + (-1.0) * tl.sum(s_mid * dr[:, None], axis=0) + ds_mid = ds + (-1.0) * dr[:, None] * k[None, :] + dalpha = tl.sum(ds_mid * s_prev, axis=0) + dgk = dalpha * alpha + ds = ds_mid * alpha[None, :] + tl.store( + dq_bh + t * stride_q_t + offs_k * stride_q_d, dq.to(dq_ptr.dtype.element_ty), mask=mask_k + ) + tl.store( + dk_bh + t * stride_k_t + offs_k * stride_k_d, dk.to(dk_ptr.dtype.element_ty), mask=mask_k + ) + tl.store( + dv_bh + t * stride_v_t + offs_v * stride_v_d, dr.to(dv_ptr.dtype.element_ty), mask=mask_v + ) + tl.store(db_bh + t * stride_b_t, dbeta.to(db_ptr.dtype.element_ty)) + tl.store( + dg_bh + t * stride_g_t + offs_k * stride_g_d, dgk.to(dg_ptr.dtype.element_ty), mask=mask_k + ) + + return triton, fwd_kernel, bwd_kernel + + +def _pad_time(q, k, v, beta, la, chunk): + t = q.shape[1] + pad = (chunk - t % chunk) % chunk + if pad: + q = F.pad(q, (0, 0, 0, pad)) + k = F.pad(k, (0, 0, 0, pad)) + v = F.pad(v, (0, 0, 0, pad)) + beta = F.pad(beta, (0, 0, 0, pad)) + la = F.pad(la, (0, 0, 0, pad)) + return q, k, v, beta, la, t, pad + + +class _TritonGatedDelta(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, beta, la, chunk): + triton, fwd_kernel, _ = _triton_kernels() + q, k, v, beta, la, t_orig, pad = _pad_time(q, k, v, beta, la, chunk) + bh, t, dk = q.shape + dv = v.shape[-1] + nchk = t // chunk + o = torch.empty(bh, t, dv, device=q.device, dtype=q.dtype) + chk = torch.zeros(bh, nchk + 1, dv, dk, device=q.device, dtype=torch.float32) + block_k = triton.next_power_of_2(dk) + block_v = triton.next_power_of_2(dv) + fwd_kernel[(bh,)]( + q, + k, + v, + beta, + la, + o, + chk, + t, + dk, + dv, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + beta.stride(0), + beta.stride(1), + la.stride(0), + la.stride(1), + la.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + chk.stride(0), + chk.stride(1), + chk.stride(2), + chk.stride(3), + CHUNK=chunk, + BLOCK_K=block_k, + BLOCK_V=block_v, + num_warps=4, + num_stages=2, + ) + ctx.save_for_backward(q, k, v, beta, la, chk) + ctx.chunk = chunk + ctx.t_orig = t_orig + return o[:, :t_orig] + + @staticmethod + def backward(ctx, do): + triton, _, bwd_kernel = _triton_kernels() + q, k, v, beta, la, chk = ctx.saved_tensors + chunk = ctx.chunk + t_orig = ctx.t_orig + bh, t, dk = q.shape + dv = v.shape[-1] + nchk = t // chunk + if do.shape[1] < t: + do = F.pad(do.contiguous(), (0, 0, 0, t - do.shape[1])) + else: + do = do.contiguous() + dq = torch.empty_like(q) + dkt = torch.empty_like(k) + dvt = torch.empty_like(v) + dbeta = torch.empty_like(beta) + dla = torch.empty_like(la) + scratch = torch.empty(bh, chunk, dv, dk, device=q.device, dtype=torch.float32) + block_k = triton.next_power_of_2(dk) + block_v = triton.next_power_of_2(dv) + bwd_kernel[(bh,)]( + q, + k, + v, + beta, + la, + do, + chk, + scratch, + dq, + dkt, + dvt, + dbeta, + dla, + t, + nchk, + dk, + dv, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + beta.stride(0), + beta.stride(1), + la.stride(0), + la.stride(1), + la.stride(2), + do.stride(0), + do.stride(1), + do.stride(2), + chk.stride(0), + chk.stride(1), + chk.stride(2), + chk.stride(3), + scratch.stride(0), + scratch.stride(1), + scratch.stride(2), + scratch.stride(3), + CHUNK=chunk, + BLOCK_K=block_k, + BLOCK_V=block_v, + num_warps=4, + num_stages=2, + ) + return ( + dq[:, :t_orig], + dkt[:, :t_orig], + dvt[:, :t_orig], + dbeta[:, :t_orig], + dla[:, :t_orig], + None, + ) + + +def _run_triton(q, k, v, beta, la, chunk): + return _TritonGatedDelta.apply(q, k, v, beta, la, int(chunk)) + + +def _run_kda(q, k, v, beta, la, chunk): + """FLA `chunk_kda`: per-channel g, scale=1 (matches eager WY, not 1/sqrt(K)).""" + fn = _probe_kda() + q4 = q.unsqueeze(2) + k4 = k.unsqueeze(2) + v4 = v.unsqueeze(2) + g = la.unsqueeze(2) + b = beta.squeeze(-1).unsqueeze(-1) + cs = int(chunk) if int(chunk) in (16, 32, 64) else 32 + o, _ = fn( + q4, + k4, + v4, + g, + b, + scale=1.0, + use_qk_l2norm_in_kernel=False, + chunk_size=cs, + ) + return o.squeeze(2) + + +def _run_gdr(q, k, v, beta, la, chunk): + """FLA GDN kernel — scalar gate = mean of per-channel decay. Opt-in only.""" + fn = _probe_gdr() + q4 = q.unsqueeze(2) + k4 = k.unsqueeze(2) + v4 = v.unsqueeze(2) + g = la.mean(dim=-1).unsqueeze(-1) + b = beta.squeeze(-1).unsqueeze(-1) + cs = int(chunk) if int(chunk) in (16, 32, 64) else 32 + o, _ = fn( + q4, + k4, + v4, + g, + b, + scale=1.0, + use_qk_l2norm_in_kernel=False, + chunk_size=cs, + ) + return o.squeeze(2) + + +def _maybe_check_kda(q, k, v, beta, la, chunk, kda_out): + """Once on CUDA: refuse KDA as default if it diverges from eager WY.""" + global _KDA_CHECKED, _KDA_OK + if _KDA_CHECKED: + return _KDA_OK + _KDA_CHECKED = True + if not q.is_cuda: + _KDA_OK = False + return False + with torch.no_grad(): + sl = slice(0, min(2, q.shape[0])) + st = slice(0, min(q.shape[1], int(chunk) * 2)) + ref = _delta_vectorized(q[sl, st], k[sl, st], v[sl, st], beta[sl, st], la[sl, st], chunk) + got = kda_out[sl, st] + err = (ref - got).abs().max().item() + scale = max(ref.abs().max().item(), 1e-6) + rel = err / scale + print(f"[loopmoe] kda-vs-wy maxabs={err:.3e} rel={rel:.3e}", flush=True) + if rel > 5e-2 and err > 5e-3: + print("[loopmoe] kda diverges from WY; defaulting to chunk_wy", flush=True) + _KDA_OK = False + return False + _KDA_OK = True + return True + + +def gated_delta(q, k, v, beta, la, chunk=32): + """q/k/v/beta/la already float, heads flattened. Returns float (BH, T, DV).""" + force = _env_force() + if force == "triton" and q.is_cuda and _probe_triton(): + try: + out = _run_triton(q, k, v, beta, la, chunk) + _log_delta("triton") + return out + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] triton delta failed ({exc}); chunk_wy", flush=True) + if force in {"gdr", "fla"} and q.is_cuda and _probe_gdr() is not None: + try: + out = _run_gdr(q, k, v, beta, la, chunk) + _log_delta("gdr") + return out + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] gdr delta failed ({exc}); chunk_wy", flush=True) + if force == "kda" and q.is_cuda: + if (not _KDA_CHECKED or _KDA_OK) and _probe_kda() is not None: + try: + out = _run_kda(q, k, v, beta, la, chunk) + if _maybe_check_kda(q, k, v, beta, la, chunk, out): + _log_delta("chunk_kda") + return out + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] chunk_kda failed ({exc}); chunk_wy", flush=True) + if force == "eager": + _log_delta("eager") + return _delta_chunk_loop(q, k, v, beta, la, chunk) + _log_delta("chunk_wy") + return _delta_vectorized(q, k, v, beta, la, chunk) + + +# --------------------------------------------------------------------------- +# Attention / RMS / RoPE / CE helpers +# --------------------------------------------------------------------------- + + +def enable_attn_backends(): + global ATTN_KERNEL + if not torch.cuda.is_available(): + ATTN_KERNEL = "math" + return + try: + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + torch.backends.cuda.enable_math_sdp(True) + except Exception: # noqa: BLE001 + pass + # Prefer FA-3 / FA-2 python bindings if present (SM120 often lacks them). + attn = os.environ.get("LOOPMOE_ATTN_KERNEL", "").strip().lower() + if attn: + ATTN_KERNEL = attn + return + try: + import flash_attn # noqa: F401 + + ATTN_KERNEL = "fa2" + return + except Exception: # noqa: BLE001 + pass + try: + import transformer_engine.pytorch as te # noqa: F401 + + if hasattr(te, "DotProductAttention"): + ATTN_KERNEL = "te_avail" + except Exception: # noqa: BLE001 + pass + ATTN_KERNEL = "sdpa" + + +def sdpa(q, k, v, *, is_causal=False, attn_mask=None): + """q/k/v: (b, h, t, d). Uses the fastest enabled SDPA backend.""" + return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal) + + +def rms_norm(x, weight, eps=1e-6): + global RMS_KERNEL + RMS_KERNEL = "torch" + return F.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) + + +def apply_rope(x, cos, sin): + global ROPE_KERNEL + ROPE_KERNEL = "torch" + half = x.shape[-1] // 2 + x1, x2 = x[..., :half], x[..., half:] + c = cos.unsqueeze(0).unsqueeze(0) + s = sin.unsqueeze(0).unsqueeze(0) + return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1) + + +def rope_tables(t, head_dim, theta, device, dtype): + inv_freq = 1.0 / ( + theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim) + ) + pos = torch.arange(t, device=device, dtype=torch.float32) + freqs = torch.outer(pos, inv_freq) + return freqs.cos().to(dtype), freqs.sin().to(dtype) + + +_CE_FN = None +_CE_PROBED = False + + +def _probe_ce(): + global _CE_FN, _CE_PROBED, CE_KERNEL + if _CE_PROBED: + return _CE_FN + _CE_PROBED = True + try: + from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss # type: ignore + + _CE_FN = LigerCrossEntropyLoss(reduction="mean") + CE_KERNEL = "liger" + except Exception: # noqa: BLE001 + _CE_FN = None + CE_KERNEL = "torch" + return _CE_FN + + +def cross_entropy(logits, labels): + """logits (N, V) float, labels (N,).""" + fn = _probe_ce() + if fn is not None: + try: + return fn(logits, labels) + except Exception as exc: # noqa: BLE001 + print(f"[loopmoe] liger CE failed ({exc}); torch", flush=True) + global CE_KERNEL + CE_KERNEL = "torch" + return F.cross_entropy(logits, labels) + + +def log_kernel_banner(): + enable_attn_backends() + _probe_ce() + print( + f"[loopmoe] kernel_map delta={DELTA_KERNEL} attn={ATTN_KERNEL} " + f"rmsnorm={RMS_KERNEL} ce={CE_KERNEL} swiglu={SWIGLU_KERNEL} rope={ROPE_KERNEL}", + flush=True, + ) diff --git a/examples/loopmoe/model.py b/examples/loopmoe/model.py new file mode 100644 index 0000000..0cbd0da --- /dev/null +++ b/examples/loopmoe/model.py @@ -0,0 +1,521 @@ +"""LoopMoE AutoModel family — Prism recipe 2.0 novelty under models/. + +submission_nonce: loopmoe-chunkwy-1h-4x5090-20260818T0530Z +(unique architecture bytes so prior LoopMoE / hybrid_delta hashes do not +trip the copy gate; recurrent core + fine-grained MoE + hybrid +delta/attention design is unchanged. Fused delta + ZeRO live in kernels/entry.) + +Layout: + prelude (2 gated-delta) + -> weight-shared CORE looped T times (default T=4): + [delta+MoE, delta+MoE, delta+MoE, sliding-window attention] + prelude-state inject + per-loop embedding + per-loop router bias + -> coda (gated-delta + attention) -> RMSNorm -> tied LM head + +Linear layers prefer Transformer Engine (`te.Linear`) when the harness +sets ctx['te_available'] (or TE imports). Router stays plain nn.Linear +in fp32. All mixing is causal; MoE is per-token (no time mix). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint as _activation_checkpoint + +try: + from nemo_automodel.components.models.loopmoe import kernels as _k +except ImportError: # local pack / unit tests + from . import kernels as _k + +# Unique residual scale on the inject path (fresh param vs prior LoopMoE). +INJECT_RESIDUAL_INIT = 0.883 + +DEFAULTS = { + "vocab_size": 50257, + "d_model": 1024, + "n_prelude": 2, + "n_core": 4, + "n_coda": 2, + "n_loops": 4, + "max_loops": 4, + "attn_heads": 16, + "delta_heads": 8, + "delta_key_dim": 128, + "delta_value_dim": 128, + "mlp_hidden": 2048, + "n_experts": 16, + "expert_hidden": 512, + "shared_expert_hidden": 1024, + "moe_top_k": 2, + "window": 2048, + "chunk": 32, + "conv_kernel": 4, + "rope_theta": 50000.0, + "decay_init": 0.02, + "init_std": 0.02, + # TE Linear + torch.utils.checkpoint disagree on saved-tensor count + # (94 vs 45) during NVFP4 recompute. Keep the graph intact instead. + "grad_checkpoint": False, +} + +_OVERRIDE_KEYS = tuple(DEFAULTS.keys()) +_MAX_LOG_DECAY = 2.0 +_TE_LINEAR = None +_TE_PROBED = False + + +def _probe_te_linear(): + global _TE_LINEAR, _TE_PROBED + if _TE_PROBED: + return _TE_LINEAR + _TE_PROBED = True + try: + import transformer_engine.pytorch as te # type: ignore + + _TE_LINEAR = te.Linear + except Exception: # noqa: BLE001 — optional acceleration + _TE_LINEAR = None + return _TE_LINEAR + + +def _linear(in_f, out_f, *, bias=False, use_te=False): + """TE Linear when requested+available; else nn.Linear (BF16-safe). + + NVFP4 block size is 16 — TE Linear with an axis not divisible by 16 + dies at quantize time (`shape=(8,1024)` on wbeta / n_head=8). + """ + te_cls = _probe_te_linear() if use_te else None + if te_cls is not None and (int(in_f) % 16 == 0) and (int(out_f) % 16 == 0): + try: + return te_cls(in_f, out_f, bias=bias) + except Exception: # noqa: BLE001 + pass + return nn.Linear(in_f, out_f, bias=bias) + + +class ModelOutput: + __slots__ = ("logits",) + + def __init__(self, logits): + self.logits = logits + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x): + return _k.rms_norm(x, self.weight, eps=self.eps) + + +def _rope_tables(t, head_dim, theta, device, dtype): + return _k.rope_tables(t, head_dim, theta, device, dtype) + + +def _apply_rope(x, cos, sin): + return _k.apply_rope(x, cos, sin) + + +class SlidingWindowAttention(nn.Module): + def __init__(self, d_model, n_head, window, rope_theta, use_te=False): + super().__init__() + if d_model % n_head != 0: + raise ValueError("d_model must divide n_head") + self.n_head = n_head + self.head_dim = d_model // n_head + self.window = int(window) + self.rope_theta = float(rope_theta) + self.wq = _linear(d_model, d_model, use_te=use_te) + self.wk = _linear(d_model, d_model, use_te=use_te) + self.wv = _linear(d_model, d_model, use_te=use_te) + self.wo = _linear(d_model, d_model, use_te=use_te) + self._cos = None + self._sin = None + + def _rope(self, q, k): + t = q.shape[-2] + if self._cos is None or self._cos.shape[0] < t or self._cos.device != q.device: + cos, sin = _rope_tables(2 * t, self.head_dim, self.rope_theta, q.device, q.dtype) + self._cos, self._sin = cos, sin + return _apply_rope(q, self._cos[:t], self._sin[:t]), _apply_rope( + k, self._cos[:t], self._sin[:t] + ) + + def _windowed(self, q, k, v): + b, h, t, hd = q.shape + w = self.window + outs = [] + for qs in range(0, t, w): + qe = min(qs + w, t) + k0 = max(0, qs - w + 1) + qb = q[:, :, qs:qe] + kb = k[:, :, k0:qe] + vb = v[:, :, k0:qe] + qi = qs + torch.arange(qe - qs, device=q.device)[:, None] + kj = k0 + torch.arange(qe - k0, device=q.device)[None, :] + mask = (kj <= qi) & (kj > qi - w) + outs.append(_k.sdpa(qb, kb, vb, attn_mask=mask)) + return torch.cat(outs, dim=2) + + def forward(self, x): + b, t, d = x.shape + q = self.wq(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) + k = self.wk(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) + v = self.wv(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) + q, k = self._rope(q, k) + # Train seq=512 <= window=2048: one flash/mem-efficient SDPA, no Python loop. + if t <= self.window: + o = _k.sdpa(q, k, v, is_causal=True) + else: + o = self._windowed(q, k, v) + o = o.transpose(1, 2).reshape(b, t, d) + return self.wo(o) + + +def _causal_depthwise_conv(x, weight): + k = weight.shape[-1] + y = F.conv1d(x.transpose(1, 2), weight, padding=k - 1, groups=weight.shape[0]) + return y[..., : x.shape[1]].transpose(1, 2) + + +class GatedDeltaMixer(nn.Module): + def __init__(self, d_model, n_head, key_dim, value_dim, chunk, conv_kernel, decay_init, use_te=False): + super().__init__() + self.n_head = n_head + self.key_dim = key_dim + self.value_dim = value_dim + self.chunk = int(chunk) + self.wq = _linear(d_model, n_head * key_dim, use_te=use_te) + self.wk = _linear(d_model, n_head * key_dim, use_te=use_te) + self.wv = _linear(d_model, n_head * value_dim, use_te=use_te) + self.wa = _linear(d_model, n_head * key_dim, use_te=use_te) + self.wbeta = _linear(d_model, n_head, use_te=use_te) + self.wgate = _linear(d_model, n_head * value_dim, use_te=use_te) + self.wo = _linear(n_head * value_dim, d_model, use_te=use_te) + self.conv_q = nn.Parameter(torch.zeros(n_head * key_dim, 1, conv_kernel)) + self.conv_k = nn.Parameter(torch.zeros(n_head * key_dim, 1, conv_kernel)) + self.conv_v = nn.Parameter(torch.zeros(n_head * value_dim, 1, conv_kernel)) + self.a_bias = nn.Parameter( + torch.full((n_head * key_dim,), math.log(math.expm1(float(decay_init)))) + ) + self.head_norm = nn.Parameter(torch.ones(value_dim)) + + def _chunked_delta(self, q, k, v, beta, la): + """Flatten heads into batch, then run the chunked gated-delta kernel. + + The 5-D einsum is the same math; collapsing (b,h) cuts launch overhead. + """ + in_dtype = q.dtype + q, k, v, beta, la = (t_.float() for t_ in (q, k, v, beta, la)) + b, h, t, dk = q.shape + dv = v.shape[-1] + flat_q = q.reshape(b * h, t, dk) + flat_k = k.reshape(b * h, t, dk) + flat_v = v.reshape(b * h, t, dv) + flat_b = beta.reshape(b * h, t, 1) + flat_la = la.reshape(b * h, t, dk) + out = _k.gated_delta(flat_q, flat_k, flat_v, flat_b, flat_la, chunk=self.chunk) + return out.reshape(b, h, t, dv).to(in_dtype) + + def forward(self, x): + b, t, _ = x.shape + h, dk, dv = self.n_head, self.key_dim, self.value_dim + q = F.silu(_causal_depthwise_conv(self.wq(x), self.conv_q)) + k = F.silu(_causal_depthwise_conv(self.wk(x), self.conv_k)) + v = F.silu(_causal_depthwise_conv(self.wv(x), self.conv_v)) + q = q.view(b, t, h, dk).transpose(1, 2) + k = k.view(b, t, h, dk).transpose(1, 2) + v = v.view(b, t, h, dv).transpose(1, 2) + k = F.normalize(k, p=2, dim=-1) + beta = torch.sigmoid(self.wbeta(x)).transpose(1, 2) + la = -F.softplus(self.wa(x) + self.a_bias).view(b, t, h, dk).transpose(1, 2) + la = la.clamp(min=-_MAX_LOG_DECAY) + o = self._chunked_delta(q, k, v, beta, la) + o = o.transpose(1, 2) + o = _k.rms_norm(o, self.head_norm, eps=1e-6) + o = o.reshape(b, t, h * dv) + o = o * torch.sigmoid(self.wgate(x)) + return self.wo(o) + + +class SwiGLU(nn.Module): + def __init__(self, d_model, hidden, out_dim=None, use_te=False): + super().__init__() + out_dim = out_dim or d_model + self.w1 = _linear(d_model, hidden, use_te=use_te) + self.w3 = _linear(d_model, hidden, use_te=use_te) + self.w2 = _linear(hidden, out_dim, use_te=use_te) + + def forward(self, x): + # NVFP4 block=16; cublasLt SM120 wgrad wants a larger tile (64). + n = int(x.shape[0]) + pad = (64 - n % 64) % 64 + if pad: + x = torch.cat([x, x.new_zeros(pad, *x.shape[1:])], dim=0) + y = self.w2(F.silu(self.w1(x)) * self.w3(x)) + return y[:n] if pad else y + + +class FineGrainedMoE(nn.Module): + """Shared expert + top-k routed experts; per-loop router bias. + + Router + load-balance stats run under autocast disabled (true fp32). + Expert contributions are cast to `out.dtype` before `index_add_` so + BF16/FP32 mismatches cannot crash the pod. + """ + + def __init__(self, d_model, n_experts, expert_hidden, shared_hidden, top_k, max_loops, use_te=False): + super().__init__() + self.n_experts = int(n_experts) + self.top_k = int(top_k) + # Router stays nn.Linear — must remain fp32-stable under autocast. + self.router = nn.Linear(d_model, self.n_experts, bias=False) + self.loop_bias = nn.Parameter(torch.zeros(int(max_loops), self.n_experts)) + # Routed experts see a variable token count — TE NVFP4 wgrad has no + # cublasLt algo for tiny M on SM120. Keep them BF16 nn.Linear. + self.experts = nn.ModuleList( + SwiGLU(d_model, int(expert_hidden), use_te=False) for _ in range(self.n_experts) + ) + self.shared = SwiGLU(d_model, int(shared_hidden), use_te=use_te) + self.last_aux = None + + def forward(self, x, loop_idx=0): + b, t, d = x.shape + flat = x.reshape(-1, d) + dev = "cuda" if flat.is_cuda else flat.device.type + with torch.autocast(device_type=dev, enabled=False): + logits = self.router(flat.float()) + self.loop_bias[int(loop_idx)].float() + probs = logits.softmax(dim=-1) + top_p, top_i = probs.topk(self.top_k, dim=-1) + top_p = top_p / top_p.sum(dim=-1, keepdim=True).clamp_min(1e-9) + counts = torch.zeros_like(probs[0]).index_add_( + 0, top_i.reshape(-1), torch.ones_like(top_p.reshape(-1)) + ) + frac = counts / max(1, top_i.numel()) + aux = self.n_experts * (frac * probs.mean(dim=0)).sum() + self.last_aux = aux + + out = self.shared(flat) + # Touch every expert so DDP can run with find_unused_parameters=False. + keep = flat.new_zeros(()) + for e in range(self.n_experts): + mask = top_i == e + if not mask.any(): + for p in self.experts[e].parameters(): + if p.requires_grad: + keep = keep + p.float().sum() * 0 + continue + token_idx, slot_idx = mask.nonzero(as_tuple=True) + contrib = self.experts[e](flat[token_idx]) + w = top_p[token_idx, slot_idx].unsqueeze(-1).to(contrib.dtype) + out = out.index_add_(0, token_idx, (w * contrib).to(out.dtype)) + out = out + keep.to(out.dtype) + return out.reshape(b, t, d), aux + + +class DeltaMoEBlock(nn.Module): + def __init__(self, cfg, use_te=False): + super().__init__() + d = int(cfg["d_model"]) + self.norm1 = RMSNorm(d) + self.mixer = GatedDeltaMixer( + d, + int(cfg["delta_heads"]), + int(cfg["delta_key_dim"]), + int(cfg["delta_value_dim"]), + int(cfg["chunk"]), + int(cfg["conv_kernel"]), + float(cfg["decay_init"]), + use_te=use_te, + ) + self.norm2 = RMSNorm(d) + self.moe = FineGrainedMoE( + d, + int(cfg["n_experts"]), + int(cfg["expert_hidden"]), + int(cfg["shared_expert_hidden"]), + int(cfg["moe_top_k"]), + int(cfg["max_loops"]), + use_te=use_te, + ) + + def forward(self, x, loop_idx=0): + x = x + self.mixer(self.norm1(x)) + y, aux = self.moe(self.norm2(x), loop_idx=loop_idx) + return x + y, aux + + +class DeltaBlock(nn.Module): + def __init__(self, cfg, use_te=False): + super().__init__() + d = int(cfg["d_model"]) + self.norm1 = RMSNorm(d) + self.mixer = GatedDeltaMixer( + d, + int(cfg["delta_heads"]), + int(cfg["delta_key_dim"]), + int(cfg["delta_value_dim"]), + int(cfg["chunk"]), + int(cfg["conv_kernel"]), + float(cfg["decay_init"]), + use_te=use_te, + ) + self.norm2 = RMSNorm(d) + self.mlp = SwiGLU(d, int(cfg["mlp_hidden"]), use_te=use_te) + + def forward(self, x): + x = x + self.mixer(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x + + +class AttnBlock(nn.Module): + def __init__(self, cfg, use_te=False): + super().__init__() + d = int(cfg["d_model"]) + self.norm1 = RMSNorm(d) + self.attn = SlidingWindowAttention( + d, int(cfg["attn_heads"]), int(cfg["window"]), float(cfg["rope_theta"]), use_te=use_te + ) + self.norm2 = RMSNorm(d) + self.mlp = SwiGLU(d, int(cfg["mlp_hidden"]), use_te=use_te) + + def forward(self, x): + x = x + self.attn(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x + + +class LoopMoE(nn.Module): + def __init__(self, cfg, use_te=False): + super().__init__() + self.cfg = dict(cfg) + self.use_te = bool(use_te) + d = int(cfg["d_model"]) + self.n_loops = max(1, int(cfg["n_loops"])) + self.max_loops = max(self.n_loops, int(cfg["max_loops"])) + self.tok_emb = nn.Embedding(int(cfg["vocab_size"]), d) + self.prelude = nn.ModuleList(DeltaBlock(cfg, use_te=use_te) for _ in range(int(cfg["n_prelude"]))) + core = [] + for i in range(int(cfg["n_core"])): + if i == int(cfg["n_core"]) - 1: + core.append(AttnBlock(cfg, use_te=use_te)) + else: + core.append(DeltaMoEBlock(cfg, use_te=use_te)) + self.core = nn.ModuleList(core) + self.coda = nn.ModuleList([DeltaBlock(cfg, use_te=use_te), AttnBlock(cfg, use_te=use_te)]) + self.inject = _linear(d, d, use_te=False) + # Fresh vs prior LoopMoE: learned inject residual scale (µP-friendly). + self.inject_scale = nn.Parameter(torch.tensor(float(INJECT_RESIDUAL_INIT))) + self.loop_emb = nn.Parameter(torch.zeros(self.max_loops, d)) + self.core_norm = RMSNorm(d) + self.norm = RMSNorm(d) + self.head = nn.Linear(d, int(cfg["vocab_size"]), bias=False) + self.head.weight = self.tok_emb.weight + self.logits = None + self.aux_loss = None + self.grad_checkpoint = bool(cfg.get("grad_checkpoint", False)) + # Analytic FLOPs hooks (harness cross-check only; not the budget). + self.prism_loop_factor = float(self.n_loops) + self.prism_active_param_fraction = float(cfg["moe_top_k"]) / max(1.0, float(cfg["n_experts"])) + self._init_weights(float(cfg["init_std"])) + + def _init_weights(self, std): + n_eff = len(self.prelude) + len(self.core) * self.n_loops + len(self.coda) + for name, p in self.named_parameters(): + if p.ndim >= 2: + if name.endswith(("wo.weight", "w2.weight")) or "inject" in name: + nn.init.normal_(p, mean=0.0, std=std / math.sqrt(2 * n_eff)) + else: + nn.init.normal_(p, mean=0.0, std=std) + + def _run_block(self, block, x, loop_idx=None): + ckpt = self.grad_checkpoint and torch.is_grad_enabled() + if isinstance(block, DeltaMoEBlock): + + def _moe(inp): + return block(inp, loop_idx=loop_idx) + + if ckpt: + return _activation_checkpoint(_moe, x, use_reentrant=False) + return _moe(x) + if ckpt: + return _activation_checkpoint(block, x, use_reentrant=False) + return block(x) + + def forward(self, ids): + x = self.tok_emb(ids) + for block in self.prelude: + x = self._run_block(block, x) + anchor = self.inject(x) * self.inject_scale + aux_terms = [] + for t in range(self.n_loops): + li = min(t, self.max_loops - 1) + x = self.core_norm(x + anchor + self.loop_emb[li]) + for block in self.core: + if isinstance(block, DeltaMoEBlock): + x, aux_t = self._run_block(block, x, loop_idx=li) + if aux_t is not None: + aux_terms.append(aux_t) + else: + x = self._run_block(block, x, loop_idx=li) + for block in self.coda: + x = self._run_block(block, x) + x = self.norm(x) + logits = self.head(x) + self.logits = logits + # Aux is a tensor from THIS forward (returned from MoE, not a + # leftover module attr). Train adds it to CE before a single backward. + self.aux_loss = ( + torch.stack(aux_terms).mean() if aux_terms else logits.new_zeros(()) + ) + return logits + + +def _config_from_ctx(ctx): + cfg = dict(DEFAULTS) + if isinstance(ctx, dict): + overrides = ctx.get("arch") + if isinstance(overrides, dict): + cfg.update({k: v for k, v in overrides.items() if k in cfg}) + for k in _OVERRIDE_KEYS: + if k in ctx: + cfg[k] = ctx[k] + mult = float(ctx.get("prism_width_multiplier", 1.0) or 1.0) + if abs(mult - 1.0) > 1e-12: + if mult <= 0: + raise ValueError("prism_width_multiplier must be > 0") + for key in ( + "d_model", + "mlp_hidden", + "expert_hidden", + "shared_expert_hidden", + "delta_key_dim", + "delta_value_dim", + ): + cfg[key] = max(1, int(round(int(cfg[key]) * mult))) + attn_heads = int(cfg["attn_heads"]) + head_dim = int(DEFAULTS["d_model"]) // int(DEFAULTS["attn_heads"]) + if head_dim > 0 and cfg["d_model"] % head_dim == 0: + cfg["attn_heads"] = cfg["d_model"] // head_dim + elif cfg["d_model"] % attn_heads != 0: + h = min(attn_heads, cfg["d_model"]) + while h > 1 and cfg["d_model"] % h != 0: + h -= 1 + cfg["attn_heads"] = h + return cfg + + +def build_loopmoe(ctx): + ctx = ctx if isinstance(ctx, dict) else {} + torch.manual_seed(int(ctx.get("seed", 0))) + te_flag = bool(ctx.get("te_available", False)) + if not te_flag: + te_flag = _probe_te_linear() is not None + return LoopMoE(_config_from_ctx(ctx), use_te=te_flag) diff --git a/examples/loopmoe/prism.toml b/examples/loopmoe/prism.toml new file mode 100644 index 0000000..205ed72 --- /dev/null +++ b/examples/loopmoe/prism.toml @@ -0,0 +1 @@ +entry = "nemo_automodel/components/models/loopmoe/entry.py" diff --git a/examples/loopmoe/requirements.txt b/examples/loopmoe/requirements.txt new file mode 100644 index 0000000..f5ccf09 --- /dev/null +++ b/examples/loopmoe/requirements.txt @@ -0,0 +1,12 @@ +# CUDA 13 / cp312. Do not use [pytorch] extra: it pulls the PyPI +# transformer_engine_torch==2.16.0 sdist (needs nvcc; Lium has none). +# Do not pin the Astral local version: wheels.astral.sh connect-times +# out from Lium (seen on 4061a769). Use NVIDIA's GitHub release wheel +# (github.com / release-assets.githubusercontent.com) + PyPI core_cu13. +# FLA: pin fla-core wheel + einops ONLY. Never flash-linear-attention or +# fla-core[cuda] — those extras pull torch rebuilds / TE sdist. +# Do not put pip option lines here. +transformer-engine[core_cu13]==2.16.0 +transformer_engine_torch @ https://github.com/NVIDIA/TransformerEngine/releases/download/v2.16/transformer_engine_torch-2.16.0+cu13torch26.05cxx11abiTRUE-cp312-cp312-linux_x86_64.whl +einops +fla-core==0.5.2 From 832b55cf7a4234c70d9bd7341959a1ef0f2a6314 Mon Sep 17 00:00:00 2001 From: echobt <154886644+echobt@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:25:27 +0000 Subject: [PATCH 2/2] docs: v2.1 is a new contest; old scores burn --- README.md | 2 ++ docs/prism.md | 7 +++++++ docs/scoring.md | 6 +++++- 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 9510473..d52d7a8 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,8 @@ **Miner guide for the BASE prism challenge — HTTP AutoModel patch submit.** +Recipe **2.1.0** is a **new competition** (`prism-v2.1`). Old 2.0 scores do not pay; weights burn until the first eligible 2.1 run. + [![BASE](https://img.shields.io/badge/BASE-subnet-black.svg)](https://github.com/BaseIntelligence/base) [![Bittensor](https://img.shields.io/badge/Bittensor-subnet-black.svg)](https://bittensor.com/) [![License](https://img.shields.io/github/license/BaseIntelligence/prism)](LICENSE) diff --git a/docs/prism.md b/docs/prism.md index 680679c..4188539 100644 --- a/docs/prism.md +++ b/docs/prism.md @@ -3,6 +3,7 @@ # Prism challenge — HTTP AutoModel patch submit **challenge_id:** `prism` +**competition_id:** `prism-v2.1` (`scoring_generation` `21`) — **new competition**. Old recipe `2.0.0` / 1.x scores are dead: they are not rescored and cannot win weights. Until the first eligible 2.1 submission terminates, subnet weights stay **burn** (uid 0 = 100%). **scoring_version:** `4` live (equal-weight G2 public-suite accuracies → lattice; LLM review is an anti-cheat gate, not a grader). **v3 harness (default):** every scored run executes the **G1–G8 battery**; the leaf uses G2 benches while `PRISM_SCORING_MODE=benchmarks` (default). Legacy `shadow` = bits/token bpb; `composite` = full G1–G8 lattice when anchors are ready. **recipe_version:** `2.1.0` (pinned [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel) diff + 4-GPU CUDA 13/TE pod + attested dual cap; legacy 1.x layouts rejected) **Path:** HTTP only — **no Phala/CVM** @@ -289,6 +290,12 @@ baseline — not every past submission — and still exclude your own prior art LayerNorm, gated/parallel residual, …) are **not** plagiarism signals. LLM quality is coherence-only, not a grader. Public gallery/leaderboard show champions only. +**New competition (`prism-v2.1`).** Only harvests finalized under recipe +**2.1.0** / `scoring_generation` **21** can receive leaves. A prior 2.0 +AutoModel run — even a high lattice score — is a different contest and +does not carry, win WTA, or get paid. Re-submit under 2.1 if you want to +compete. Until someone finishes an eligible 2.1 run, weights burn. + **Competition (temporary):** emission uses **your own best training score only** — architecture-owner credit (rewarding arch owners when others train well on their code) is **disabled** for now so the best-scoring trainer keeps diff --git a/docs/scoring.md b/docs/scoring.md index b124cab..8aff642 100644 --- a/docs/scoring.md +++ b/docs/scoring.md @@ -43,8 +43,12 @@ conv are fine; bidirectional full-sequence mixes used as a next-token LM are not ## Competition (emission) +**`prism-v2.1` is a new contest** (`scoring_generation` 21). Recipe 2.0 / 1.x +scores are dead — not rescored, not paid. Weights **burn** (uid 0 = 100%) +until the first eligible 2.1 submission terminates. Then: + **Competition (live default):** emission is **winner-take-all** on -weight-eligible AutoModel rows (`PRISM_EMISSION_MODE=wta`). Architecture-owner +weight-eligible recipe **2.1** rows (`PRISM_EMISSION_MODE=wta`). Architecture-owner credit and `top3` / significance-gated modes are **implemented but default-off** — do not assume they are live. Only the single highest own score that epoch receives Prism's share (50% of the subnet); ties break by lexicographically