emo integration - #841
emo integration#841AkshitaB wants to merge 63 commits into
Conversation
💡 Codex ReviewOLMo-core/src/olmo_core/train/callbacks/batch_size_scheduler.py Lines 175 to 177 in f8d6d10  Rebuild the PP schedule when changing the batch size
When an OLMo-core/src/olmo_core/launch/beaker.py Lines 1047 to 1049 in f8d6d10  Preserve an explicitly selected stable image
When a preset is combined with an explicit ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
💡 Codex ReviewOLMo-core/src/olmo_core/ops/moe.py Line 362 in 679b8b0 Return one gradient per async all-to-all input
When a training path using OLMo-core/src/olmo_core/train/train_module/transformer/ddp_train_module.py Lines 1961 to 1963 in 679b8b0 Make pipeline evaluation return output on every rank
With pipeline parallelism enabled, ranks that do not own the final stage return OLMo-core/src/olmo_core/launch/beaker.py Lines 1041 to 1042 in 679b8b0 Let explicit setup flags replace preset steps
When a caller combines a preset with ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
💡 Codex ReviewOLMo-core/src/olmo_core/distributed/parallel/pipeline_parallel.py Lines 796 to 803 in c9dca56 Preserve support for standard pipeline schedules
This rejects every pre-existing pipeline schedule, including the default ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
💡 Codex Review Handle standard schedules before iterating pipeline outputs
When Remove the unsupported pipeline-evaluation kwarg
Whenever an ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
💡 Codex ReviewOLMo-core/src/olmo_core/train/train_module/transformer/ddp_train_module.py Lines 1227 to 1230 in 8abfa62 Refresh FP8 caches after eval-only checkpoint loads
When OLMo-core/src/olmo_core/train/train_module/transformer/ddp_train_module.py Lines 1363 to 1365 in 8abfa62 Reject incomplete eval-only checkpoint loads
If an eval checkpoint lacks any parameter expected by the configured model—for example after a config mismatch or parameter rename—this silently skips that parameter while loading all other weights. Since eval-only model construction initializes parameters before loading, the missing parameter remains randomly initialized and evaluation can complete with corrupted results; collect missing names and fail the load instead of accepting any nonempty partial state. OLMo-core/src/olmo_core/nn/attention/__init__.py Lines 710 to 714 in 8abfa62 Shard attention-sink logits under tensor parallelism
With OLMo-core/src/olmo_core/nn/attention/backend.py Lines 334 to 342 in 8abfa62 Build the sink mask after Ulysses redistribution
For full attention with sinks and Ulysses context parallelism, this mask is built from the pre-redistribution local sequence lengths ( ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
💡 Codex Review Restore optimizer state even when checkpoint saves fail
If the distributed checkpoint write raises because of an I/O or upload failure, control never reaches the subsequent ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
💡 Codex ReviewOLMo-core/src/olmo_core/utils.py Lines 559 to 560 in 44adac8 Guard the stream lookup patch on supported older Torch
When running with PyTorch 2.6, which ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
Co-authored-by: Jacob Morrison <jacob1morrison@gmail.com>
Adds inference-time expert restriction to Olmo3MoeRouter for the EMo modularity metric: a non-persistent allowed_experts buffer with set/clear methods, applied to the router scores immediately before the ordinary token-level top-k. Masking scores rather than logits matches the training-time EMo document-pool operator, which selects on masked scores but weights with the unmasked ones. The mask is evaluation state rather than model configuration, so it stays out of the state dict and a single checkpoint can be scored repeatedly under different subsets. Also adds get_moe_routers() to collect routers by layer index, skipping dense layers. Tests cover all-experts parity as an exact no-op, exclusion of masked experts, restoration after clearing, rejection of malformed and undersized masks, and non-persistence. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
EMo derives per-token `segment_ids` from token IDs and EOS positions, but only the first pipeline stage receives token IDs — later stages receive hidden states. Both the model and the pipeline schedule rejected the combination outright. The batch is already available on every pipeline rank, and keyword arguments (unlike positional ones) reach every stage, so segment IDs need no transport. They are derived once per rank and passed through as a side input. - `emo_block_indices` / `emo_eos_token_id` on `Transformer` replace the per-forward block scan. Both resolve lazily, after the pipeline split, so a stage sees only the blocks it owns. - `_prepare_inputs` accepts caller-supplied `segment_ids`, deriving them from token IDs only as the non-pipelined fallback. It has to be an explicit pop — unrecognized kwargs are otherwise dropped silently. - Both EMo+PP guards removed (`_prepare_inputs` and `apply_pp`). The adjacent context-parallelism guard stays. - `segment_ids`, `doc_lens`, and `max_doc_lens` are now split into microbatches. `max_doc_lens` is a per-instance Python list, so it is sliced rather than passed through the tensor branch. Intra-document masking was already broken under PP for this reason, independent of EMo. - Every batch-leading value splits on one set of shared, equal boundaries instead of independent `tensor_split` calls. - Uneven microbatches are rejected. Stages size their P2P buffers from a single floor-divided microbatch shape (`prepare_step` → `example_p2p_tensor`), so larger leading chunks would overrun the receiver — a buffer shape mismatch, not just inconsistent metadata. The check sits after the active microbatch count is selected, so training, evaluation, and reduced dry runs each validate against the count they actually run with. - Batch size is inferred from any batch-leading kwarg, not just `labels`, so a stage receiving `segment_ids` without `labels` still works. `run_pipeline_eval` padding now covers batch-leading lists and tuples, which a list-valued `max_doc_lens` would otherwise be left short by. The reduced-microbatch dry-run path already handled these. - Two-batch overlap still rejects EMo separately: segment IDs are not split between lanes. - Context parallelism is unchanged and still rejected. - Cross-*rank* disagreement on the EOS token ID is unvalidated; catching it needs a collective. Disagreement across the model parts a single rank holds is caught. `pytest src/test/` passes apart from pre-existing failures (GCS credential tests, and the reshard suite, both confirmed identical on a clean tree). New coverage: - Model-side: EMo block/EOS discovery, derivation vs. caller-supplied precedence, the PP guard, shape validation, and routing segment IDs to blocks given a hidden-state input. - Splitter: alignment with `input_ids`, packed-document metadata, uneven-batch rejection, batch-size inference on later stages. - Evaluation padding for lists, tuples, and tensors. - A multi-GPU execution test asserting each stage receives the correct microbatch's segment IDs, including a rank that gets no token IDs at all. **This one has not been run** — it requires GPUs. A second commit temporarily adds this base branch to the workflow triggers, mirroring the existing `akshitab/moe-v2-core` entries. Revert before merging the base branch into `main`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
44adac8 to
82fb4a7
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 82fb4a70f4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self.disable_mxfp8_expert_anchor_grads() | ||
|
|
||
| dp_group = dense_process_group if dense_process_group is not None else dp_mesh.get_group() | ||
| load_balancing_group = dp_mesh.get_group() |
There was a problem hiding this comment.
When context parallelism is enabled, each CP rank owns only a sequence shard, but this group reduces counts only across the dp_mesh. The loss is subsequently gradient-reduced across the combined DP+CP group, so each CP shard weights its local router scores using a different partial count vector; averaging those products is not equivalent to computing the loss from global assignment counts and scores. Use the supplied combined dense_process_group (or otherwise reduce across both DP and CP) so global_load_balancing=True trains against the intended global distribution.
Useful? React with 👍 / 👎.
| if already_interpolated: | ||
| paths.append(f"{base_dir}{path}") |
There was a problem hiding this comment.
For Dolma3p5_14t, this branch ignores the tokenizer argument and always returns files tokenized with allenai/dolma2-tokenizer. Selecting the new mix with any other tokenizer configuration can therefore either feed out-of-range token IDs to the embedding or silently train on IDs with the wrong vocabulary semantics. Validate that the requested tokenizer is Dolma2-compatible instead of accepting every tokenizer.
Useful? React with 👍 / 👎.
| if getattr(hf_config, "model_type", None) == "olmo3moe": | ||
| roundtrip_state = convert_state_from_hf(hf_config, hf_state_dict, model_type="olmo3moe") |
There was a problem hiding this comment.
Every olmo3moe export now reconstructs a complete OLMo state while both the original state and converted HF state are still live. The reverse converter stacks and concatenates the per-expert tensors, so this is another full-model allocation rather than a lightweight view; for the tens-of-billions-parameter models supported by this change it adds hundreds of gigabytes of peak host memory and can make an otherwise viable conversion OOM before writing anything. Validate mappings incrementally or make this expensive round-trip check opt-in.
Useful? React with 👍 / 👎.
4 participants
No description provided.