Skip to content

start script - #7

Merged
samsja merged 9 commits into
mainfrom
start_script
Feb 26, 2025
Merged

start script#7
samsja merged 9 commits into
mainfrom
start_script

Conversation

@samsja

@samsja samsja commented Feb 26, 2025 •
edited
Loading

Copy link
Copy Markdown
Member

example:

 uv run src/zeroband/rl_launcher.py --n_gpus 8  --train @ configs/training/150M/H100.toml --inference @ configs/inference/debug.toml --inference.name_model 150M --inference.max_samples 10000

output

04:28:45 [INFO] [RL_LAUNCHER] [Rank 0] start rl training with 4 GPUs, 4. Total: 8
04:28:45 [INFO] [RL_LAUNCHER] [Rank 0] train_gpus_ids: [4, 5, 6, 7]
04:28:45 [INFO] [RL_LAUNCHER] [Rank 0] inference_gpus_ids: [0, 1, 2, 3]

on 8xH100

using 4 gpus for training and 4 gpus for inference

uv run src/zeroband/rl_launcher.py --n_gpus 8  --train @ configs/training/150M/H100.toml --inference @ configs/inference/debug.toml --inference.name_model 150M --inference.max_samples 10000 --train.data.fake

Screenshot from 2025-02-25 21-11-49

@samsja
add rl launcher script
a8ecc1d 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

add setup stuff inside main train

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

add rl launcher

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

fix device placements

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

udpate logger

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
add rl launcher
70f6aa0 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
updat ereadme
70bcc3a 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
samsja marked this pull request as draft February 26, 2025 02:42
@samsja
add tp to vllm config
e382e7a 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
refactor order
d1b4e08 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
fix local rank bug
060b48c 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
add tp support
1ac7b12 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja
samsja marked this pull request as ready for review February 26, 2025 05:13
apaz-cli
apaz-cli approved these changes Feb 26, 2025
View reviewed changes

@apaz-cli apaz-cli left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great. Is there maybe a way to launch it that's less verbose? Maybe it could default to use all your GPUs?

Comment thread tests/test_train.py Outdated
process = subprocess.Popen(cmd)
result = process.wait()
if result != 0:
pytest.fail(f"Process {result} failed {result}")

@apaz-cli apaz-cli Feb 26, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Assume you meant to print something else?

@samsja samsja Feb 26, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

uh no that's what I wanted to print haha

@Jackmin801 Jackmin801 Feb 26, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

u sure?

@samsja samsja Feb 26, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh realized the problem lol

@samsja samsja Feb 26, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will blame @Jackmin801 for not catching this when doing prime https://github.com/PrimeIntellect-ai/prime/blob/d57965b04574262815a174246707afd9615eed0c/tests/test_torchrun/test_train.py#L78

@samsja samsja Feb 26, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

@Jackmin801 Jackmin801 Feb 26, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That line was added in a PR without review: PrimeIntellect-ai/prime#3
PrimeIntellect-ai/prime@f34f9c9

Blame reverted

@samsja
add auto gpu
9613396 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
@samsja

samsja commented Feb 26, 2025

Copy link
Copy Markdown
Member Author

Looks great. Is there maybe a way to launch it that's less verbose? Maybe it could default to use all your GPUs?

yeah great idea, I added an option to get auto gpu

@samsja
samsja requested a review from apaz-cli February 26, 2025 06:58
@samsja
fix double results
812f7a9 
Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>
Jackmin801
Jackmin801 approved these changes Feb 26, 2025
View reviewed changes

@Jackmin801 Jackmin801 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm! abit hacky in some places but agi will refactor them in 2026 so no worries

Comment thread tests/test_train.py Outdated
process = subprocess.Popen(cmd)
result = process.wait()
if result != 0:
pytest.fail(f"Process {result} failed {result}")

@Jackmin801 Jackmin801 Feb 26, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

u sure?

@samsja
samsja merged commit 11f9d64 into main Feb 26, 2025
samsja pushed a commit that referenced this pull request Mar 30, 2026
@rasdani
Merge pull request #7 from PrimeIntellect-ai/daniel/update-12-28
59ae09c 
Daniel/update 12 28
snimu added a commit that referenced this pull request Jun 4, 2026
@snimu @claude
fix(losses): review batch — validators, echo export, doc notes
5364546 
- 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>
snimu added a commit that referenced this pull request Jun 4, 2026
@snimu @claude
fix(losses): 2nd review — echo-only filtering, drop sft/opd presets, …
afd7129 
…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>
@cursor cursor Bot mentioned this pull request Jun 9, 2026
Closed
1 task
snimu added a commit that referenced this pull request Jun 10, 2026
@snimu @claude
fix(losses): pack reference_logprobs via the optional-field helper (b…
f2591fb 
…ug review #7)

Packing used a bespoke inline branch for reference_logprobs that initialized an
empty list when the bin had none and a later sample did (mis-aligning, not
back-filling existing_len zeros) and had no branch for the reverse case. Use the
existing _extend_optional_token_field helper, exactly as overlays do — back-fills
the bin on first appearance and pads when a later sample lacks it.

Latent today (reference_logprobs is OPD-only and batches don't pack mixed
training_mode, so mixed presence is unreachable), but reference scoring is going
general in the OPD/reference follow-up, so get the foundation right here.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@mikasenghaas
mikasenghaas deleted the start_script branch August 5, 2026 04:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@apaz-cli apaz-cli Awaiting requested review from apaz-cli
1 more reviewer
@Jackmin801 Jackmin801 Jackmin801 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@samsja @apaz-cli @Jackmin801