Skip to content

fix(configs): validate inference weight_broadcast type against orchestrator#2848

Open
EazyReal wants to merge 1 commit into
PrimeIntellect-ai:mainfrom
EazyReal:fix/weight-broadcast-3way-validation
Open

fix(configs): validate inference weight_broadcast type against orchestrator#2848
EazyReal wants to merge 1 commit into
PrimeIntellect-ai:mainfrom
EazyReal:fix/weight-broadcast-3way-validation

Conversation

@EazyReal

@EazyReal EazyReal commented Jun 21, 2026

Copy link
Copy Markdown
Contributor

What

validate_shared_weight_broadcast checks that the trainer, orchestrator, and (optional) inference weight-broadcast transports agree. It used a chained comparison:

trainer.weight_broadcast.type != orchestrator.weight_broadcast.type != inference.weight_broadcast.type

Python evaluates this as (trainer != orchestrator) and (orchestrator != inference), so inference is never compared against trainer. When trainer and orchestrator agree but inference differs (e.g. filesystem / filesystem / nccl), the check passes and the mismatch only surfaces as a broken weight sync at runtime.

This is reachable whenever the shared [weight_broadcast] is unset (the default): auto_setup_weight_broadcast only unifies the three sub-configs when the shared field is set, otherwise it passes the user's per-component values straight to this validator.

Why / fix

Compare each component against the orchestrator independently, which transitively requires all three to match. Both original error messages are preserved.

if trainer.weight_broadcast.type != orchestrator.weight_broadcast.type:
    raise ValueError("Trainer ... and orchestrator ... are not the same ...")
if inference is not None and inference.weight_broadcast.type != orchestrator.weight_broadcast.type:
    raise ValueError("Inference ... and orchestrator ... are not the same ...")

Test

tests/unit/test_configs.py:

  • test_weight_broadcast_inference_mismatch_raises(filesystem, filesystem, nccl) now raises (fails on main: "DID NOT RAISE").
  • test_weight_broadcast_all_matching_passes — all-matching config still validates.

All 111 tests in test_configs.py pass; ruff format/ruff check clean.

…trator

validate_shared_weight_broadcast used a chained comparison
`trainer.type != orchestrator.type != inference.type`, which Python
evaluates as `(trainer != orchestrator) and (orchestrator != inference)`.
Inference was therefore never compared against trainer, so a config where
trainer and orchestrator agree but inference differs (e.g.
filesystem/filesystem/nccl) passed validation and broke weight sync at
runtime. This is reachable whenever the shared [weight_broadcast] is unset,
since auto_setup_weight_broadcast otherwise unifies all three components.

Check each component against the orchestrator independently, which
transitively requires all three to match.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@EazyReal EazyReal force-pushed the fix/weight-broadcast-3way-validation branch from 22d317c to dfb98f4 Compare June 30, 2026 08:41
@EazyReal

Copy link
Copy Markdown
Contributor Author

@mikasenghaas @samsja Could I get a quick review on this config validation fix? It catches inference weight-broadcast type mismatches before runtime weight sync breaks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant