Configure attention and feed-forward norms independently#687
Conversation
Split `TransformerBlockConfig.layer_norm` into separate `attention_norm` and `feed_forward_norm` fields so the pre-attention and pre-feed-forward layer norms can be configured independently. Previously a single `layer_norm` config was used to build both norms. For backwards compatibility, `layer_norm` is retained as a deprecated `InitVar`: when provided it is copied into both `attention_norm` and `feed_forward_norm`, and specifying both the legacy and split fields raises an `OLMoConfigurationError`. Since it is an `InitVar` it is not stored or serialized, so existing on-disk configs that use `layer_norm` load and round-trip cleanly into the new split fields. Block constructors now accept `attention_norm` and `feed_forward_norm` in place of a single `layer_norm`, and the `llama_like` / `gemma3_like` factories populate both fields. Behavior is unchanged when the same norm config is used for both. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The merged-config test indexed the serialized block config at `block["layer_norm"]`, which no longer exists now that the block config serializes `attention_norm` / `feed_forward_norm`. Point it at `attention_norm`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
@codex review |
![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: 371dd68701
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
… separately PeriNormTransformerBlock builds post_attention_norm from attention_norm and post_feed_forward_norm from feed_forward_norm, so the two extra norms can have different parameter counts when configured independently. Count one of each instead of doubling feed_forward_norm. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The layer_norm InitVar is not serialized, so overrides targeting it through config.merge() (e.g. model.block.layer_norm.eps=1e-6) silently fail to apply. Note the limitation in the CHANGELOG and field docstring, directing users to override attention_norm/feed_forward_norm instead. This matches the existing behavior of the deprecated `attention` InitVar. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
One follow-up concern: since this PR now allows Today |
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>
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>
2 participants
Summary
Split
TransformerBlockConfig.layer_norminto separateattention_normandfeed_forward_normfields, so the pre-attention and pre-feed-forward layer norms can be configured independently. Previously a singlelayer_normconfig was used to build both.Backwards compatibility
layer_normis kept as a deprecatedInitVar:attention_normandfeed_forward_norm.layer_normand either split field raisesOLMoConfigurationError.InitVar, it is never stored or serialized — existing on-disk configs that uselayer_normload and round-trip cleanly into the new split fields.So this is a non-breaking change: configs and scripts that pass
layer_norm=continue to work unchanged.Changes
TransformerBlockConfig: newattention_norm/feed_forward_normfields;layer_normdemoted to a deprecatedInitVarwith mirroring in__post_init__;num_paramsreads the split fields.block.py: block constructors (TransformerBlock,LayerNormScaledTransformerBlock,PeriNormTransformerBlock,MoETransformerBlock,MoEHybridTransformerBlockBase) now takeattention_normandfeed_forward_norminstead of a singlelayer_norm. No behavior change when the same norm is used for both.llama_likeandgemma3_likepopulate both fields.Tests
config_test.py: legacylayer_norm=back-compat mirroring, independent split norms, the conflict error, and factory population of the split fields.block_test.py(builds block modules directly).layer_normon disk) passes — validating the back-compat path end to end.pytest src/test/nn/transformer/passes;make checks(isort/black/ruff/mypy) clean.🤖 Generated with Claude Code