Use hf model by samsja · Pull Request #1 · PrimeIntellect-ai/prime-rl

samsja · 2025-02-19T02:04:37Z

No description provided.

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

Co-authored-by: samsja <55492238+samsja@users.noreply.github.com>

* formatask init * add task def * formatask configs * fix formatting?? * fix + overhaul difflib reward * complete reward function definition * autoclean rollouts #1 Co-authored-by: samsja <55492238+samsja@users.noreply.github.com> * autoclean rollouts #2 Co-authored-by: samsja <55492238+samsja@users.noreply.github.com> * expanded multiepoch dataset * ruff fix --------- Co-authored-by: kalomaze <66376113+kalomaze@users.noreply.github.com> Co-authored-by: kalomaze <juicebox22245@gmail.com>

* Next prod resume config * Go to 52k seq len * Old run name

[fix] output trajectory.json on local

* bump verifier * add mika new env * add all envs * update envs * reduce sandbox resources and change sandbox label for `mini-swe-agent-plus` (#2) * update envs * bump prime-rl to main * update prime-rl * update prime-rl * update env --------- Co-authored-by: Mika Senghaas <mail@mikasenghaas.de> Co-authored-by: rasdani <73563550+rasdani@users.noreply.github.com>

Hypothesis: orchestrator's /pause returns before inference CUDA work finishes draining. Trainer starts RDMA-writing while inference is still using that memory, causing non-deterministic partial corruption that compounds over steps. Fix: - Inference: cuda.sync before SPG barrier (drain in-flight work). - SPG barrier #1: signal "inference is quiescent, trainer can write". - Trainer: blocks at matching barrier, THEN posts writes + drains. - SPG barrier #2: signal "writes done" (existing). - Inference: cuda.sync after #2, then MLA absorb. Reverted iter14 knobs (flush_every=1, enforce_eager) to isolate the effect of the pre-write barrier alone. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The previous version of `test_buffer_no_positive_resampling_save_load_round_trip` asserted `pass_rate_stats[hash(eb.examples[1])]` after save/load, but the test setup never actually exercised example_id=1: `make_rollouts(buffer, "env_a", [1], ...)` indexes positionally into `list(eb.examples.values())`, and after update PrimeIntellect-ai#1 evicted example 0, that positional 1 resolved to example_id=2. So the saved stats covered (h(ex0), h(ex2)) and the assertion looked up h(ex1) — `KeyError`. Save/load itself is correct: hashes are stable across a fresh Buffer over the same dataset, and `move_saved_pool` matches by hash. Verified on baker's internal prime-rl fork — they use the equivalent schema (list of excluded example dicts + hash-keyed stats dict) and recover example_ids from the new dataset via the same hash_lookup pattern. We leave our storage schema as-is. Test changes: - Use `[0]` for the second update (positional 0 = example_id=1 once 0 is evicted), and snapshot the actual example dicts at update time so the round-trip assertion can hash them independently of post-eviction layout. - Add an explicit comment documenting `make_rollouts`'s positional-index API to prevent the same off-by-one in future tests. - Add `test_buffer_no_positive_resampling_save_load_hashes_match_new_dataset` as a regression on the underlying invariant: hashing `eb.examples[i]` from a freshly-constructed Buffer must equal the save-time hash. If Buffer init ever mutates example dicts in a way that touches `hash_keys`, this test fails fast. Both tests + the existing 13 buffer tests pass green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The five SLURM Jinja templates upstream from prime-rl ran two commands that directly contradict the recipe invariants documented in README.md / AGENTS.md / SCALERL.md: 1. `uv sync --all-extras` (vs the documented `--extra all`). `--all-extras` enumerates every named extra and pulls in `flash-attn-3`, whose wheel ships Hopper sm_90 kernels only and crashes on B200 with "no kernel image available." Both shipped ScaleRL configs target B200. Switch to the aggregate `[all]` extra and add `--locked` so per-job syncs don't silently re-resolve the seven git-pinned deps against upstream HEAD. 2. The multi-node trainer leg unconditionally exported `HF_HUB_OFFLINE=1` while the single-node `scripts/scalerl_smoke.sbatch` exports `HF_HUB_OFFLINE=0`. Operators on clusters where `$HOME`/`HF_HOME` isn't NFS/EFS-mounted to compute nodes hit an opaque "model files not found" failure with no doc pointer. Honor any operator-set value (`HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"`) so `.env` or the parent shell can flip it to `0` for fresh-pull runs. The single-node smoke (`scripts/scalerl_smoke.sbatch`) escapes PrimeIntellect-ai#1 because it bypasses `uv sync` entirely and uses the pre-built `.venv`, but the multi-node path described in `docs/SCALERL.md` Multi-node section was directly broken by these template defaults.

Updates after read-only inspection of the cluster: - TTT plan + cheat-sheet now carry the actual cluster facts: H200 topology, /beegfs storage, vLLM 0.21.0 pin, weight broadcast=nccl, Qwen3-4B model cache, production memory pressure (optim_cpu_offload). - gate_up_proj layerwise fix is already in main (PR #2482); historical failure mode noted but no longer needs reproduction. Live SFT-on-tool runs validate the SFT path behaviorally — smoke task #1 is implicit. - New hard invariant from daniel/gptoss-lora-nan-repro HANDOVER: every chunk snapshot gets a fresh lora_int_id, never reused. LRU eviction handles cleanup. Codified in the cheat-sheet's "things that burned us" and as plan invariant #11. - Multi-LoRA infra is production-hardened via hosted training; TTT risk surface narrowed to the chunked-snapshot churn pattern, not the multi-LoRA machinery itself. New files: - docs/ttt-probes.md: four single-purpose benchmarks the cluster session should run before writing Phase A code — snapshot churn, learner forward+backward, multi-LoRA forward overhead at n_adapters ∈ {1,8,64,256,1024}, /beegfs write throughput. Each has pass criteria that gate later phases. - docs/ttt-config-draft.toml: starter TOML mirroring the production Forth qwen-rl.toml shape, with the TTT additions on top. Drops as `qwen-rl-ttt.toml` next to qwen-rl.toml in configs/private when the cluster session moves it over.

- enabled_losses=None now validated as the full term list, so >1 echo term per env is caught at config time instead of at rollout time. [review #6] - loss_overrides keys validated against `losses`; non-echo overrides rejected. [#7] - warn (don't fail) when prompt-role echo is configured with renderer=None (MITO), where prompt_attribution is unavailable so it would silently no-op. [#8] - token_export: add echo_mask/echo_weight columns + export sequences trained only via echo (gate on loss_mask OR echo_mask). [#9] - doc notes: echo CE uses the rollout temperature (scale alpha to compensate, kept as-is); negative alpha is intentional (suppresses tokens). [#1, #10] - tests for the new config validators. Deferred to a follow-up pass (per the review): full per-sample primary routing / rl-disable [#2b] + the <=1-primary validation it enables [#5], and the multi-run losses fingerprint [#3]. Not run locally; ruff + py_compile clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…sft warning - Drop sft/opd presets from the losses DSL: they dispatch to fixed cores by training_mode and are not loss-list terms, so listing them (or naming them in enabled_losses) was a silent no-op footgun. `losses` now holds rl/custom + echo only. [review #2] - Gate the zero-advantage filter on the env's rl primary being active (ZeroAdvantageFilter.primary_active, wired to train_sink._primary_enabled), so echo-only / rl-disabled envs (where advantage is irrelevant — e.g. group_size=1) aren't dropped on zero advantage. Default always-active → unchanged without echo. [review #1] - The MITO prompt-role-echo warning also fires for training_mode='sft' (which forces renderer=None after the warning's original check). [review #7] - tests for sft-type rejection + the filter gate. Left as-is: trainer.loss migration (intentional break), backfill prompt-role echo (best-effort, known), docs DPPO clipping (pre-existing main doc bug). Not run locally; ruff + py_compile clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…ht resolver - #1 reserve loss-term names {sft, opd} (any term) and `rl` (non-primary), so an overlay can't silently overwrite a training_mode dispatch core or the rl primary in the trainer registry. - #5 resolve the primary's advantage weight orchestrator-side: scale the per-token advantage by the advantage-weight's tau in process_group and drop adv_tau from the dppo_kl core / RLLossConfig. Now *any* primary core (dppo_kl or custom) gets the resolved advantage × tau — no per-core special-case. Bit-identical for the default tau=1.0. - #4 overlay trainability = non-None AND non-zero, so a zero-weight overlay (e.g. advantage-weighted with zero advantage) no longer keeps an otherwise-empty batch alive past the empty-batch guard. - #6 custom overlay weight resolver is group-aware: it now receives `WeightInputs{sample, rollouts}` (the full GRPO group) instead of a lone sample, so it can compute group-relative weights. - #10 document the overlay_mask/overlay_weight token-export columns (schema v2) in the configs skill. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…s to rl mode Two related bugs in primary λ/reduce/hooks handling (bug review #1 + #2): #1 — a custom-core primary (grpo advantage + `loss = custom`) is dispatched under "rl" via is_primary, but train.py picked the λ/reduce/hooks term by `loss.type == "dppo_kl"`, so a custom primary silently lost its lambda_weight / reduce / hooks. Now the tuning term is selected by is_primary (the dppo_kl lookup stays only for token-export's DPPO threshold annotations, which need that core). #2 — compute_loss applied primary_lambda/reduce/hooks to whatever core training_mode dispatched, including the fixed sft/opd cores, so a configured rl term's knobs could leak into sft/opd training. Gate them to training_mode == "rl" (sft/opd: λ=1, global-mean reduce, no hooks) — matching the "sft/opd cores are fixed" design. Default configs are unaffected (the leak only bit non-default rl tuning under sft/opd). Test: sft ignores a non-unit primary_lambda + a loss-zeroing primary hook. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

## Summary ### Model and checkpoint contract - Replace the separate standard, GPT-OSS, and Nemotron expert/router implementations with one canonical MoE tree. - Store routed expert weights as separate `gate_proj`, `up_proj`, and `down_proj` parameters. Non-gated experts omit `gate_proj`. - Nest every custom router at `mlp.router.gate`, routed experts at `mlp.experts`, and optional shared experts inside `mlp`. - Store the single selection-only router bias at `mlp.router.selection_bias`; model-specific checkpoint names are converted only at the HF boundary. - Keep shared expert projections canonical. Qwen3.5 alone nests its scalar sigmoid gate at `mlp.shared_expert.output_gate`, with conversion mapping the standalone Hugging Face key at the boundary. - Use the same activation protocol and `FeedForward` implementation for dense MLPs, shared experts, and grouped experts. - Make shared-expert construction exclusively parent-owned: model layers pass an explicit `FeedForward`/subclass or `None`, while `MoEArgs` and `MoE` never infer or construct one. - Keep existing HF source layouts at the conversion boundary. GPT-OSS alone adds the structural split/join required for its interleaved source tensor. - Require `.prime-v1` in reusable `<conversion_dir>/prime` caches so legacy Prime layouts fail before loading. ### MoE runtime contract - Use one routed execution sequence for every backend: dispatch tokens, run canonical grouped experts, combine tokens. - Select grouped expert compute through explicit BF16, DeepGEMM FP8, or MXFP8 implementations. `MoEArgs` remains architecture-only. - Select local/Torch or DeepEP token dispatch once during model setup. Dispatchers own routing metadata, communication, padding, chunking, and synchronization. - Keep `TokenDispatcher` as the structural interface while a shared base owns the canonical run/synchronize sequence; local, Torch, and DeepEP implementations own only their dispatch/combine details. - Keep DeepEP combine events, concatenate streams, and flattened dispatch state owned by each dispatcher; chunking uses the equal input shape shared by EP ranks and does not synthesize empty chunks. - Ignore dispatch configuration when expert parallelism resolves to one rank; local dispatch remains active and DeepEP has no optimizer side effects. - Keep grouped expert execution directly in `GroupedExperts.forward` and router confidence accounting directly in `TokenChoiceTopKRouter.forward`. - Shard every canonical expert parameter with one backend-independent `Shard(0)` placement. - Keep dense linear quantization, routed-expert compute, and EP transport independent. - Remove the TorchTitan module hooks, expert decorator, dynamically attached process group, backend attributes, tensor-subclass interception, and process-global MXFP8 patches from the reachable path. - Vendor the narrow MXFP8 grouped-GEMM and Torch EP transport implementation in [prime-kernels PR #1](PrimeIntellect-ai/prime-kernels#1), preserving the pinned torchao revision and BSD license. ## Public contracts Checkpoint tree: - `mlp.router.gate`: `nn.Linear` - `mlp.router.selection_bias`: optional persistent selection-only buffer - `mlp.experts.gate_proj`: stacked `nn.Parameter` for gated experts, absent for non-gated experts - `mlp.experts.up_proj`: stacked `nn.Parameter` - `mlp.experts.down_proj`: stacked `nn.Parameter` - `mlp.shared_expert.{gate_proj,up_proj,down_proj}`: `FeedForward` projections when present - `mlp.shared_expert.output_gate`: Qwen3.5-only scalar `nn.Linear` - expert matrices: `[num_experts, output_features, input_features]` Runtime configuration: ```toml [trainer.model.moe.compute] type = "bf16" # bf16 | deepgemm_fp8 | mxfp8 [trainer.model.moe.dispatch] type = "torch" # torch | deepep transport = "bf16" # bf16 | mxfp8 for torch dispatch ``` DeepEP owns its tuning fields: ```toml [trainer.model.moe.dispatch] type = "deepep" num_sms = 20 token_chunk_size = 4096 ``` This is intentionally breaking. `enable_grouped_gemm`, `enable_a2a`, `ep_comm_backend`, `deepep_num_sms`, and `deepep_token_chunk_size` are removed without aliases. ## Validation completed - Foundation conversion, detector, marker, state-loading, checkpoint, NIXL graph, and wire-format tests passed. - Focused CPU configuration and MoE contracts: 143 passed. - Focused MoE runtime configuration parametrizations after the EP=1 no-op change: 6 passed. - Affected custom-model CUDA suite: 24 passed, 3 skipped across AFMoE, GLM, GPT-OSS, Nemotron, Qwen3/Qwen3.5, and state loading. - Final H200 model suite after execution-path cleanup: 81 passed, 3 skipped (Slurm 2597). - Torch EP on two H200 GPUs passed with BF16 and DeepGEMM FP8, top-k 2, both score locations, shared experts, gated/non-gated/bias variants, uneven routing, a rank receiving zero routed tokens, and input/router/expert gradients. - Final Torch EP BF16 and DeepGEMM FP8 matrix passed after the cleanup (Slurm 2596). - DeepEP on two H200 GPUs passed with BF16 and DeepGEMM FP8, chunking on/off, both score locations, shared experts, a rank receiving zero routed tokens, and gradients. - Final DeepEP BF16 and DeepGEMM FP8 matrix passed after the cleanup (Slurm 2595). - Parent-owned shared-expert construction: 16 focused CPU tests and 64 H200 model tests passed, with 3 expected skips (Slurm 2600). - Full-graph MoE and shared `FeedForward` compile validation passed after the ownership cleanup (Slurm 2601). - Local DeepGEMM FP8 passed across all supported topology, activation, and bias combinations. - `torch.compile(fullgraph=True)` forward/backward passed for the full MoE and shared `FeedForward` across all supported topology/activation combinations. - Single-rank DCP and Torch EP + FSDP DCP model/optimizer resume passed. - The `mxfp8_moe` wheel build includes its public module, manifest, and license; the module imports against the pinned torchao build. - Unsupported H200 MXFP8 selection fails before model mutation. - Review audit: every queued DeepEP combine event is now synchronized, including chunked execution (two-H200 Slurm 2604). - Dispatcher review cleanup passed on two H200s: DeepEP BF16/DeepGEMM with chunking and a zero-receive rank (Slurm 2847), plus Torch EP topology/bias/zero-receive coverage (Slurm 2848). - The exact `reverse_text_moe` integration passed end-to-end on two H200s with an isolated conversion cache (Slurm 2607). - Conversion ranks now make the create decision on master and validate `.prime-v1` only after the final all-rank barrier. - Canonical router selection bias: 12 CPU MoE/conversion tests passed; CUDA checkpoint loading, Nemotron, and GLM model tests passed, with the existing GLM BF16 gradient comparison passing on isolated rerun after one threshold-edge miss (Slurm 2851/2852). - Ruff, Ruff formatting, compileall, and `git diff --check` passed. ## Deliberate exclusions - MultiLoRA is not adapted to the new runtime and may fail. - Fused-MoE trainer integration is removed and will be reconsidered only after the canonical runtime is finalized. - Routed-expert selective activation checkpointing is removed; the remaining selective targets are unchanged. ## Remaining before marking ready - Run numerical MXFP8 grouped-GEMM and MXFP8 Torch EP transport forward/backward on SM100 hardware. The available Slurm fleet currently has only H200 nodes. - Merge prime-kernels PR #1, publish wheels containing `mxfp8_moe`, and update the `pyproject.toml`/lock wheel pins. The currently pinned v0.8.0 wheel does not contain this module.  --- > [!NOTE] > **High Risk** > This changes the MoE forward/backward path, expert-parallel dispatch, checkpoint key layout, and trainer TOML in breaking ways across all custom MoE models. > > **Overview** > **Breaking trainer config:** MoE is configured via `[trainer.model.moe.compute]` (`bf16`, `deepgemm_fp8`, `mxfp8`) and `[trainer.model.moe.dispatch]` (`torch` with `transport`, or `deepep` with `num_sms` / `token_chunk_size`). Dense `[trainer.model.quantization]` no longer drives expert GEMMs or EP transport. Removed flags include `ep_comm_backend`, `deepep_*`, `moe_use_grouped_mm`, `moe_fused_kernel`, and quantization `enable_grouped_gemm` / `enable_a2a`. > > **Unified MoE execution:** All custom models share one `MoE` path—`TokenChoiceTopKRouter`, stacked `GroupedExperts` (`gate_proj` / `up_proj` / `down_proj`), and a pluggable `TokenDispatcher` (`LocalTokenDispatcher`, `TorchTokenDispatcher`, `MXFP8TorchTokenDispatcher`, `DeepEPTokenDispatcher`). `configure_moe_runtime` wires grouped GEMM backends and dispatch at setup; EP weight sharding is only `ExpertWeightParallel`. DeepEP dispatch/combine, grouped-GEMM permutation, and chunk pipelining live in the new dispatcher layer instead of ad hoc `MoE` hooks. > > **Checkpoint and model surface:** PrimeRL checkpoints rename routed weights from `w1/w2/w3` to `gate/down/up_proj`, move selection bias to `mlp.router.selection_bias`, and add GPT-OSS HF↔prime conversion. `LatentMoE` / Nemotron-specific expert classes are folded into the canonical tree; dense layers use `FeedForward`. Reusable `prime` conversion caches must include a `.prime-v1` marker. > > **Removed / docs:** Fused `flash_moe` trainer integration and `bench_fused_moe.py` are dropped; `routed_experts` selective activation checkpointing is removed. Docs and examples (e.g. GLM-5) are updated for the new tables. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit ffd5be7. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>

The stream's first episode is now #1, so the last of n reads as n, and that number is also what addresses it - the index carries it, and a reader seeks by the offset beside it. The stable mask is back beside the mismatch: the mismatch tooltip names both logprobs and the estimator, while the mask reads in probabilities, since eps is a probability distance. The sort control names itself the way the filter button does. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

samsja added 2 commits February 19, 2025 01:20

refactor: use hf transformer :(

6eb517d

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

fix tests

8e59add

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

samsja merged commit 72a1649 into main Feb 19, 2025

kalomaze added a commit that referenced this pull request Jun 12, 2025

autoclean rollouts #1

f1bdcc6

Co-authored-by: samsja <55492238+samsja@users.noreply.github.com>

faustomilletari mentioned this pull request Sep 18, 2025

Unable to run multi-node training #941

Closed

samsja pushed a commit that referenced this pull request Nov 7, 2025

Next prod resume config (#1)

9b91979

* Next prod resume config * Go to 52k seq len * Old run name

leonardtang pushed a commit to haizelabs/prime-rl that referenced this pull request Jan 1, 2026

Merge pull request PrimeIntellect-ai#1 from haizelabs/dev/hotfix

68afb06

[fix] output trajectory.json on local

snimu added a commit that referenced this pull request Jun 5, 2026

test(losses): reserved-name collision rejection (#1)

8b11286

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

mikasenghaas deleted the use-hf-model branch August 5, 2026 04:27

samsja mentioned this pull request Aug 6, 2026

exp(dynamo): researcher-ready native sidecar integration on current main #3213

Open

S1ro1 mentioned this pull request Aug 27, 2026

refactor(moe): unify runtime and checkpoint layout #3411

Merged

Navigation Menu

Uh oh!

Use hf model - #1

Use hf model#1
samsja merged 2 commits into
mainfrom
use-hf-model

samsja commented Feb 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

samsja commented Feb 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant