Port OLMoDDP transformer block + model#764
Conversation
![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: db85cdacdf
ℹ️ 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".
| rope.warmup_cache(max_seq_len, device) | ||
|
|
||
| if self.lm_head is not None and not self.tie_word_embeddings: | ||
| if self.lm_head is not None: |
There was a problem hiding this comment.
When tie_word_embeddings=True, self.to_empty() has already allocated fresh storage for the embedding and LM-head parameters, so the model must call _tie_weights() again before returning. This branch now runs for tied models and there is no subsequent re-tie, leaving the LM head as an independent output matrix despite the config and silently training/saving the wrong architecture.
Useful? React with 👍 / 👎.
| rowwise_fp8_cfg = self.rowwise_fp8 | ||
| use_rowwise_fp8 = ( | ||
| rowwise_fp8_cfg is not None | ||
| and rowwise_fp8_cfg.enabled | ||
| and mlp_inp.device.type == "cuda" | ||
| ) | ||
| use_rowwise_fp8 = rowwise_fp8_cfg is not None and rowwise_fp8_cfg.enabled |
There was a problem hiding this comment.
For shared-only blocks with rowwise_fp8.enabled=True on CPU, this now sets use_rowwise_fp8 solely from the config and immediately calls assert_runtime_supported(), which raises because rowwise FP8 requires CUDA. The previous device guard fell back to the normal shared expert path on CPU, so CPU smoke tests or CPU materialized/eval runs with an FP8 config now fail before reaching the dense fallback.
Useful? React with 👍 / 👎.
Take olmo-ddp's nn/ddp/block.py and model.py wholesale, then reconcile with core: pull te/cpu_offload forward, lazy-import the deferred rowwise_wave/deepep_v2 backends, make nvtx optional, and drop olmo-ddp cruft (dead selective-checkpoint machinery, a shadowed apply_ep stub, scattered imports, unused locals). Re-apply core cleanups the wholesale checkout dropped: keep the single-layer_norm block config (split only in the constructor), guard CUDA-event allocation behind torch.cuda.is_available() with the _tbo1 names the merged EP modules expect, and restore the block-level num_flops_per_token (routing attention through num_flops_per_token instead of a config-level flops_per_seq core lacks). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <taotianhua@outlook.com>
Pull forward olmo-ddp's context-parallel sequence-sharding helper onto the base Transformer. It shards input_ids/labels via the existing _cp_load_balancer.batch_shard and returns the original sequence length so pipeline stages can rebuild RoPE buffers consistently. Self-contained (core already has the CP load balancer); full CP train wiring lands with the DDP train module. The rest of olmo-ddp's shared transformer/block.py + model.py diff is intentionally not taken: the shared-block norm split (#687) is rejected in favor of core's single layer_norm, weight-tying is kept, and the v1 MoEHybrid overlap isn't on the OLMoDDP path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <taotianhua@outlook.com>
db85cda to
23af3c2
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 23af3c26d0
ℹ️ 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".
| "When TBO and EP no-sync are enabled, EP shared_slots must be >= 2 " | ||
| f"(block={moe_block.block_idx}, got {moe_block.ep.shared_slots})." | ||
| ) | ||
| if moe_block.ep.path != ExpertParallelPath.rowwise_nvshmem: |
There was a problem hiding this comment.
With two_batch_overlap=True, this rejects every routed block whose EP path is sync_1d or no_sync_1d. Those paths are still built by the checked-in TBO tests (src/test/nn/moe/v2/tbo_test.py and src/test/nn/moe/v2/ep_no_sync_tbo_1d_test.py) and were supported by the old combined_forward_ep_tbo/ep_no_sync_tbo_1d dispatch, so applying EP to those configs now fails before forward instead of exercising the existing TBO implementation.
Useful? React with 👍 / 👎.
| block_inp = x | ||
| del x | ||
|
|
||
| attn_res_out = self._res_norm_attn(block_inp, **kwargs) |
There was a problem hiding this comment.
For shared-only DDP blocks, checkpoint_attn=True no longer does anything because this path now calls _res_norm_attn directly instead of _checkpointed_res_norm_attn. Configs using shared-only blocks as dense layers with attention checkpointing will retain the full attention activations and can OOM compared with the previous behavior.
Useful? React with 👍 / 👎.
Re-apply moe-v2-core behaviors dropped by the wholesale port: - re-tie word embeddings after to_empty() in init_weights (a tied model was silently left with an independent LM head after materialization); - fall back to the dense shared-expert path on CPU when rowwise_fp8 is enabled (the CUDA-only assert broke CPU smoke/eval/materialize runs); - honor checkpoint_attn for shared-only blocks by routing through _checkpointed_res_norm_attn instead of _res_norm_attn. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <taotianhua@outlook.com>
olmo-ddp's OLMoDDP forward_tbo is rowwise_nvshmem-only, superseding moe-v2-core's 1d two-batch-overlap (ep_sync_tbo / ep_no_sync_tbo_1d). Remove the now-orphaned 1d TBO modules and their tests; the model's TBO requirement check already restricts to the rowwise_nvshmem path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <taotianhua@outlook.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d5a59437b6
ℹ️ 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".
| max_local_microbatch_size: Optional[int] = None, | ||
| device: Optional[torch.device] = None, | ||
| world_mesh: Optional[Dict[str, Optional[DeviceMesh]]] = None, | ||
| world_mesh: Dict[str, Optional[DeviceMesh]], |
There was a problem hiding this comment.
When OLMoDDPModel.init_weights() is called through the public config/model API without pipeline or expert-parallel meshes, this new required keyword raises TypeError before the body can run; for example the existing DDP smoke path builds the model and calls model.init_weights(device=...). world_mesh is only used under self.pp_enabled/self.ep_enabled, so standalone/non-parallel initialization should keep the base Transformer.init_weights() behavior by making this argument optional.
Useful? React with 👍 / 👎.
world_mesh was a required keyword, so the non-parallel init path (model.init_weights(device=...)) raised TypeError before running. It's only used under pipeline/expert parallelism, so default it to None and assert presence in those branches. Also move the misplaced docstring above the local import so it's a real docstring. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <taotianhua@outlook.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 78d0960d3d
ℹ️ 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".
| cpu_offload_handler = CpuOffloadHandler( | ||
| num_offload_group=num_layers, | ||
| num_model_group=model_layers, | ||
| tensor_need_offloading_checker=tensor_need_offloading_checker_activations, | ||
| ) |
There was a problem hiding this comment.
When callers request the documented disabled path (enabled=False) or run CPU-only smoke/materialization code, this still constructs CpuOffloadHandler, whose initializer immediately allocates torch.cuda.Stream() objects. That makes the supposed no-op helper raise on CPU-only processes before it can return (nullcontext(), None); construct the handler only inside the if enabled: branch.
Useful? React with 👍 / 👎.
get_cpu_offload_context(enabled=False) constructed CpuOffloadHandler, whose __init__ allocates torch.cuda.Stream(), so the documented no-op/disabled path raised on CPU-only processes. Build the handler only when enabled; return (nullcontext(), None) otherwise. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <taotianhua@outlook.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d1e830f6ef
ℹ️ 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".
| "When TBO and EP no-sync are enabled, EP shared_slots must be >= 2 " | ||
| f"(block={moe_block.block_idx}, got {moe_block.ep.shared_slots})." | ||
| ) | ||
| if moe_block.ep.path != ExpertParallelPath.rowwise_nvshmem: |
There was a problem hiding this comment.
Fresh evidence after the prior TBO comment: src/test/nn/moe/v2/ep_sync_tbo_test.py:67-69 still builds the checked-in TBO comparison with ExpertParallelPath.sync_1d, and test_sync_ep_tbo_matches_standard_forward runs that model with two_batch_overlap=True. This guard raises in forward_tbo() before the test can execute, so either the synced TBO path needs to remain accepted here or the remaining sync_1d TBO test/config needs to be removed or updated; otherwise the current multi-GPU test suite fails for the existing sync_1d TBO coverage.
Useful? React with 👍 / 👎.
…ad, restore config docstrings - Thread per-block context-parallel RoPE shards into the routed rowwise-nvshmem TBO loop, matching the dense/non-TBO paths so TBO + CP blocks receive their pos_sin/pos_cos/freqs_cis shards. - Remove the unused, untested TE cpu-offload module (te/cpu_offload.py, te/__init__.py) and the dead commented offload hooks in OLMoDDPModel; the feature was never wired up. - Restore the class docstring and per-attribute docstrings on OLMoDDPTransformerBlockConfig lost in the wholesale replacement. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com>
ep_sync_tbo_test.py built the TBO model with ExpertParallelPath.sync_1d, but _check_tbo_requirements() now rejects every EP path except rowwise_nvshmem, so the test would fail at the path guard on a GPU runner. The supported rowwise-NVSHMEM TBO surface is already covered by ep_no_sync_tbo_rowwise_test.py, and the other legacy TBO tests were already removed with the 1d-TBO retirement. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com>
tbo_state.py held only SyncedTboPendingContext, which nothing constructs or consumes now that _tbo_last_step() is narrowed to the no-sync rowwise pending context. Its only reference was a smoke assertion that existed solely to keep the dead type alive. Drop the module and the assertion so the retired synchronous-TBO surface is removed consistently. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com>
1 participant
Ports the fused OLMoDDP transformer block and model, bringing the MoE-v2 block/model up to their full feature set.
What's here
nn/ddp/block.py+model.py: theOLMoDDPTransformerBlock/OLMoDDPModelcompute, EP dispatch, rowwise-FP8, TBO, and symmetric-memory buffer management.nn/moe/v2/te/cpu_offload.py: Transformer-Engine CPU-offload helpers (currently unwired — offload is disabled by default; wiring lands with the TE work).nn/moe/v2/block.py: convenience re-exports of the router/expert configs.Transformer.prepare_cp_sequence_inputs: context-parallel sequence-sharding helper on the base model (shardsinput_ids/labelsvia the existing CP load balancer and returns the original sequence length for RoPE rebuild).Notes
rowwise_wave/deepep_v2backends are imported lazily at their dispatch sites and remain rejected byExpertParallelConfig.layer_normfield (split into attention/feed-forward norms in the constructor); CUDA events are allocated only when CUDA is available;num_flops_per_tokenis implemented on the block.src/test/nn/ddpsuite; the CP + train-module integration test lands with the DDP train module.🤖 Generated with Claude Code