Skip to content

Port OLMoDDP transformer block + model #764

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AkshitaB merged 9 commits into from
Jul 14, 2026
Merged

Port OLMoDDP transformer block + model #764

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
b979f40
Port OLMoDDP block + model wholesale from olmo-ddp
AkshitaB Jul 13, 2026
23af3c2
Add Transformer.prepare_cp_sequence_inputs (CP precursor)
AkshitaB Jul 13, 2026
3328e60
Restore lost core behaviors in OLMoDDP block/model (Codex)
AkshitaB Jul 14, 2026
d5a5943
Retire moe-v2-core's 1d TBO in favor of olmo-ddp's rowwise TBO
AkshitaB Jul 14, 2026
78d0960
Make OLMoDDPModel.init_weights world_mesh optional
AkshitaB Jul 14, 2026
d1e830f
Don't allocate CUDA streams in disabled cpu-offload path
AkshitaB Jul 14, 2026
c33ce51
Address PR-C review: thread CP kwargs into TBO, drop unused cpu-offlo…
AkshitaB Jul 14, 2026
5c1ce67
Remove retired synchronous-TBO test
AkshitaB Jul 14, 2026
61d03c6
Remove orphaned synchronous-TBO state module
AkshitaB Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
514 changes: 269 additions & 245 deletions src/olmo_core/nn/ddp/block.py
View file
Open in desktop

Large diffs are not rendered by default.

737 changes: 339 additions & 398 deletions src/olmo_core/nn/ddp/model.py
View file
Open in desktop

Large diffs are not rendered by default.

6 changes: 6 additions & 0 deletions src/olmo_core/nn/moe/v2/block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@
OLMoDDPTransformerBlock,
OLMoDDPTransformerBlockConfig,
)
from olmo_core.nn.moe.v2.routed_experts import RoutedExpertsConfig
from olmo_core.nn.moe.v2.router import MoERouterConfigV2
from olmo_core.nn.moe.v2.shared_experts import SharedExpertsConfig

MoEFusedV2TransformerBlock = OLMoDDPTransformerBlock
MoEFusedV2TransformerBlockConfig = OLMoDDPTransformerBlockConfig
Expand All @@ -17,4 +20,7 @@
"OLMoDDPTransformerBlockConfig",
"MoEFusedV2TransformerBlock",
"MoEFusedV2TransformerBlockConfig",
"MoERouterConfigV2",
"RoutedExpertsConfig",
"SharedExpertsConfig",
]
Loading
Loading