feat: Add EXAONE 4.0 model bridge (LG AI Research) - #2532
Conversation
Signed-off-by: Bias92 <pewpewplay315@gmail.com>
2df4d8d to
fc798ba
Compare
📝 WalkthroughWalkthroughIntroduces EXAONE 4.0 model support to Megatron's bridge infrastructure by adding new provider classes, transformer layer specifications, and bidirectional configuration conversion between HuggingFace and Megatron formats. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~30 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (3)
src/megatron/bridge/models/exaone/exaone4_provider.py (1)
58-58: Unusedlogger— dead code.
logger = logging.getLogger(__name__)is defined but there are no logging calls in this file. Remove it (and theimport logging) unless you intend to add structured logging, in which case useprint_rank_0per the project convention.♻️ Proposed fix
-import logging from dataclasses import dataclass ... - -logger = logging.getLogger(__name__)As per coding guidelines: "Use
print_rank_0for logging in model bridge to avoid duplicate output across ranks."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/megatron/bridge/models/exaone/exaone4_provider.py` at line 58, Remove the unused logger and its import: delete the standalone "logger = logging.getLogger(__name__)" and the "import logging" statement in exaone4_provider.py; if you need any runtime messages in functions/classes in this file, replace uses of logger with the project convention print_rank_0(...) instead (search for "logger" and the top-level import to update accordingly).src/megatron/bridge/models/__init__.py (1)
296-298:__all__entries are out of alphabetical model-family order.The EXAONE entries are inserted between the Qwen3 and DeepSeek blocks. Moving them before the DeepSeek group (or after all the Qwen entries) would match the alphabetical ordering followed by the rest of
__all__.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/megatron/bridge/models/__init__.py` around lines 296 - 298, The __all__ list is not alphabetically ordered: the entries "Exaone4Bridge", "Exaone4ModelProvider", and "Exaone4ModelProvider1P2B" are inserted between Qwen3 and DeepSeek blocks; move these three EXAONE entries so they appear in alphabetical model-family order (i.e., place them before the DeepSeek group or immediately after the Qwen entries) to restore consistent ordering in __all__.src/megatron/bridge/models/exaone/exaone4_bridge.py (1)
37-53: Dead code:loggeris defined but never used.
loggingis imported andlogger = logging.getLogger(__name__)is created, but there are no logging calls in this file. Remove both if logging is not needed; otherwise add relevantprint_rank_0calls for key bridge events per the project convention for model bridges.♻️ Proposed fix
-import logging - import torch ... - -logger = logging.getLogger(__name__)As per coding guidelines: "Use
print_rank_0for logging in model bridge to avoid duplicate output across ranks."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/megatron/bridge/models/exaone/exaone4_bridge.py` around lines 37 - 53, The file defines logger and imports logging but never uses them; either remove the unused import and the logger = logging.getLogger(__name__) line, or replace any intended log calls with the project convention print_rank_0. Specifically, remove the unused logging import and the logger symbol from exaone4_bridge.py if no runtime messages are needed, or add print_rank_0 calls (not logger) inside relevant bridge functions/classes such as MegatronModelBridge-related methods, Exaone4ModelProvider hooks, or PreTrainedCausalLM integration points to report key bridge events.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/megatron/bridge/models/__init__.py`:
- Around line 28-32: The bridge implementation file must be renamed to follow
project conventions: move/rename the current exaone4_bridge.py into
model_bridge.py and update any imports to import Exaone4Bridge,
Exaone4ModelProvider, and Exaone4ModelProvider1P2B from model_bridge instead of
exaone4_bridge; also fix the import ordering in __init__.py so the Exaone*
import block appears in strict alphabetical order among other model imports
(place the Exaone imports in the correct spot relative to Deepseek and others)
to match the established import ordering convention.
In `@src/megatron/bridge/models/exaone/exaone4_bridge.py`:
- Line 115: The code uses single-quoted string literals in the getattr call for
rotary_base; change the single quotes to double quotes so the line reads using
double-quoted strings for the attribute names (update the getattr(hf_config,
"rope_theta", getattr(hf_config, "rotary_base", 1000000.0)) invocation),
preserving the same logic and default value and leaving variable names like
rotary_base and hf_config unchanged.
- Around line 150-157: The code unconditionally emits hard-coded RoPE scaling
keys in the hf_config block (low_freq_factor, high_freq_factor,
original_max_position_embeddings) causing lossy megatron_to_hf_config
round-trips; update the logic to persist and reuse the full rope_scaling dict on
the provider (e.g., add a provider.rope_scaling_dict or store
provider.rope_scaling as the complete dict) and change the hf_config assignment
in exaone4_bridge (where hf_config["rope_scaling"] is set) to read and copy
values from that persisted dict instead of using fixed literals; also update the
provider_bridge/megatron_to_hf_config code path to populate
provider.rope_scaling from incoming model config so the round-trip restores
original low_freq_factor, high_freq_factor, and
original_max_position_embeddings.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py`:
- Line 35: Remove the unused Optional import from the typing import in
exaone4_provider.py: update the import line that currently reads "from typing
import Callable, Optional, Union" to drop Optional so it becomes "from typing
import Callable, Union" (or equivalent import ordering used in the file) to
satisfy static analysis and avoid the unused-import warning; ensure no other
occurrences of Optional are referenced in functions or type hints such as any in
ExaOne4Provider-related methods before committing.
- Line 175: Change the type annotation for transformer_layer_spec to use PEP 604
union syntax: replace Union[ModuleSpec, Callable[["GPTModelProvider"],
ModuleSpec]] with ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec]
(symbol: transformer_layer_spec, types: ModuleSpec, Callable, GPTModelProvider,
exaone4_layer_spec); also remove any now-unused Union import from typing.
- Around line 91-94: TERowParallelLinearLayerNorm.forward currently calls
super().forward(x) -> (output, bias) then applies self.post_layernorm(output)
before bias is added; add a guard to prevent silently incorrect behavior by
checking the returned bias: either assert bias is None with a clear message
(e.g., raise AssertionError("TERowParallelLinearLayerNorm expects no deferred
bias; set add_bias_linear=False or handle bias before layernorm")), or if you
want to support bias, add bias to output before calling self.post_layernorm
(i.e., output = output + bias) and return (normalized_output, None); make the
change inside TERowParallelLinearLayerNorm.forward and reference
super().forward, self.post_layernorm, and the bias variable.
- Line 102: The function exaone4_layer_spec currently declares a parameter
config but never uses it; either remove the unused parameter or clearly mark it
as an intentional placeholder: if you intend future config-dependent specs,
retain the parameter and add a comment above exaone4_layer_spec explaining that
config is intentionally unused for now (e.g., "placeholder for future
config-dependent layer specs such as 32B hybrid attention"), otherwise change
the signature to def exaone4_layer_spec() -> ModuleSpec and update all call
sites that pass a config to call the new no-arg function.
- Line 126: Import DotProductAttention from
megatron.core.transformer.dot_product_attention and explicitly assign it to the
core_attention argument instead of leaving core_attention=None in ExaOne4
provider; update the import list in
src/megatron/bridge/models/exaone/exaone4_provider.py (matching the pattern used
in gpt_provider.py) and set core_attention=DotProductAttention where the ExaOne4
model is constructed so the provider uses the same explicit attention class as
Gemma2/Gemma3/MCore-based providers.
---
Nitpick comments:
In `@src/megatron/bridge/models/__init__.py`:
- Around line 296-298: The __all__ list is not alphabetically ordered: the
entries "Exaone4Bridge", "Exaone4ModelProvider", and "Exaone4ModelProvider1P2B"
are inserted between Qwen3 and DeepSeek blocks; move these three EXAONE entries
so they appear in alphabetical model-family order (i.e., place them before the
DeepSeek group or immediately after the Qwen entries) to restore consistent
ordering in __all__.
In `@src/megatron/bridge/models/exaone/exaone4_bridge.py`:
- Around line 37-53: The file defines logger and imports logging but never uses
them; either remove the unused import and the logger =
logging.getLogger(__name__) line, or replace any intended log calls with the
project convention print_rank_0. Specifically, remove the unused logging import
and the logger symbol from exaone4_bridge.py if no runtime messages are needed,
or add print_rank_0 calls (not logger) inside relevant bridge functions/classes
such as MegatronModelBridge-related methods, Exaone4ModelProvider hooks, or
PreTrainedCausalLM integration points to report key bridge events.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py`:
- Line 58: Remove the unused logger and its import: delete the standalone
"logger = logging.getLogger(__name__)" and the "import logging" statement in
exaone4_provider.py; if you need any runtime messages in functions/classes in
this file, replace uses of logger with the project convention print_rank_0(...)
instead (search for "logger" and the top-level import to update accordingly).
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/megatron/bridge/models/__init__.pysrc/megatron/bridge/models/exaone/__init__.pysrc/megatron/bridge/models/exaone/exaone4_bridge.pysrc/megatron/bridge/models/exaone/exaone4_provider.py
| from megatron.bridge.models.exaone import ( | ||
| Exaone4Bridge, | ||
| Exaone4ModelProvider, | ||
| Exaone4ModelProvider1P2B, | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Bridge file should be named model_bridge.py per project conventions.
The bridge implementation is in exaone4_bridge.py, but project conventions (established for all existing model bridges) require the file to be model_bridge.py. This import transitively surfaces that naming deviation.
Additionally, the exaone import block is inserted at lines 28–32 — before the deepseek block at line 33 — breaking the alphabetical ordering used across all other model imports in this file.
♻️ Suggested import ordering fix
-from megatron.bridge.models.exaone import (
- Exaone4Bridge,
- Exaone4ModelProvider,
- Exaone4ModelProvider1P2B,
-)
from megatron.bridge.models.deepseek import (
...
)
+from megatron.bridge.models.exaone import (
+ Exaone4Bridge,
+ Exaone4ModelProvider,
+ Exaone4ModelProvider1P2B,
+)Based on learnings: "Create model bridge implementation in 'model_bridge.py' when adding new model bridges" (applies to src/megatron/bridge/models/*/model_bridge.py).
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| from megatron.bridge.models.exaone import ( | |
| Exaone4Bridge, | |
| Exaone4ModelProvider, | |
| Exaone4ModelProvider1P2B, | |
| ) | |
| from megatron.bridge.models.deepseek import ( | |
| Deepseek4Bridge, | |
| Deepseek4ModelProvider, | |
| Deepseek4ModelProvider1P2B, | |
| ) | |
| from megatron.bridge.models.exaone import ( | |
| Exaone4Bridge, | |
| Exaone4ModelProvider, | |
| Exaone4ModelProvider1P2B, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/__init__.py` around lines 28 - 32, The bridge
implementation file must be renamed to follow project conventions: move/rename
the current exaone4_bridge.py into model_bridge.py and update any imports to
import Exaone4Bridge, Exaone4ModelProvider, and Exaone4ModelProvider1P2B from
model_bridge instead of exaone4_bridge; also fix the import ordering in
__init__.py so the Exaone* import block appears in strict alphabetical order
among other model imports (place the Exaone imports in the correct spot relative
to Deepseek and others) to match the established import ordering convention.
| seq_length=hf_config.max_position_embeddings, | ||
| init_method_std=hf_config.initializer_range, | ||
| layernorm_epsilon=hf_config.rms_norm_eps, | ||
| rotary_base=getattr(hf_config, 'rope_theta', getattr(hf_config, 'rotary_base', 1000000.0)), |
There was a problem hiding this comment.
Single-quote string literals violate the double-quote coding guideline.
Line 115 uses 'rope_theta' and 'rotary_base' with single quotes.
♻️ Proposed fix
- rotary_base=getattr(hf_config, 'rope_theta', getattr(hf_config, 'rotary_base', 1000000.0)),
+ rotary_base=getattr(hf_config, "rope_theta", getattr(hf_config, "rotary_base", 1000000.0)),As per coding guidelines: "Use double quotes for strings (matching ruff formatter configuration)."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| rotary_base=getattr(hf_config, 'rope_theta', getattr(hf_config, 'rotary_base', 1000000.0)), | |
| rotary_base=getattr(hf_config, "rope_theta", getattr(hf_config, "rotary_base", 1000000.0)), |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_bridge.py` at line 115, The code
uses single-quoted string literals in the getattr call for rotary_base; change
the single quotes to double quotes so the line reads using double-quoted strings
for the attribute names (update the getattr(hf_config, "rope_theta",
getattr(hf_config, "rotary_base", 1000000.0)) invocation), preserving the same
logic and default value and leaving variable names like rotary_base and
hf_config unchanged.
| if provider.rope_scaling: | ||
| hf_config["rope_scaling"] = { | ||
| "rope_type": "llama3", | ||
| "factor": provider.rope_scaling_factor, | ||
| "low_freq_factor": 1.0, | ||
| "high_freq_factor": 4.0, | ||
| "original_max_position_embeddings": 8192, | ||
| } |
There was a problem hiding this comment.
Hard-coded RoPE scaling constants cause lossy megatron_to_hf_config round-trips.
low_freq_factor, high_freq_factor, and original_max_position_embeddings are not stored on the provider, so they are unconditionally emitted as 1.0, 4.0, and 8192 regardless of the source model's actual values. Any checkpoint with different values will produce a mis-configured HF config on the return path. Consider persisting these on the provider, or at minimum reading them back from a stored dict.
♻️ Suggested approach: store the full rope_scaling dict on the provider
In provider_bridge:
+ # Store full rope_scaling dict for lossless round-trip
if hf_rope_scaling is not None and hf_rope_scaling.get("rope_type") == "llama3":
provider.rope_scaling = True
provider.rope_scaling_factor = hf_rope_scaling.get("factor", 16.0)
+ provider.rope_scaling_original_max_position_embeddings = hf_rope_scaling.get(
+ "original_max_position_embeddings", 8192
+ )
+ provider.rope_scaling_low_freq_factor = hf_rope_scaling.get("low_freq_factor", 1.0)
+ provider.rope_scaling_high_freq_factor = hf_rope_scaling.get("high_freq_factor", 4.0)In megatron_to_hf_config:
if provider.rope_scaling:
hf_config["rope_scaling"] = {
"rope_type": "llama3",
"factor": provider.rope_scaling_factor,
- "low_freq_factor": 1.0,
- "high_freq_factor": 4.0,
- "original_max_position_embeddings": 8192,
+ "low_freq_factor": provider.rope_scaling_low_freq_factor,
+ "high_freq_factor": provider.rope_scaling_high_freq_factor,
+ "original_max_position_embeddings": provider.rope_scaling_original_max_position_embeddings,
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_bridge.py` around lines 150 - 157,
The code unconditionally emits hard-coded RoPE scaling keys in the hf_config
block (low_freq_factor, high_freq_factor, original_max_position_embeddings)
causing lossy megatron_to_hf_config round-trips; update the logic to persist and
reuse the full rope_scaling dict on the provider (e.g., add a
provider.rope_scaling_dict or store provider.rope_scaling as the complete dict)
and change the hf_config assignment in exaone4_bridge (where
hf_config["rope_scaling"] is set) to read and copy values from that persisted
dict instead of using fixed literals; also update the
provider_bridge/megatron_to_hf_config code path to populate
provider.rope_scaling from incoming model config so the round-trip restores
original low_freq_factor, high_freq_factor, and
original_max_position_embeddings.
|
|
||
| import logging | ||
| from dataclasses import dataclass | ||
| from typing import Callable, Optional, Union |
There was a problem hiding this comment.
Remove unused Optional import (confirmed by static analysis).
Optional is imported but never referenced in this file.
♻️ Proposed fix
-from typing import Callable, Optional, Union
+from typing import Callable, Union📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| from typing import Callable, Optional, Union | |
| from typing import Callable, Union |
🧰 Tools
🪛 Flake8 (7.3.0)
[error] 35-35: 'typing.Optional' imported but unused
(F401)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py` at line 35, Remove the
unused Optional import from the typing import in exaone4_provider.py: update the
import line that currently reads "from typing import Callable, Optional, Union"
to drop Optional so it becomes "from typing import Callable, Union" (or
equivalent import ordering used in the file) to satisfy static analysis and
avoid the unused-import warning; ensure no other occurrences of Optional are
referenced in functions or type hints such as any in ExaOne4Provider-related
methods before committing.
| def forward(self, x): | ||
| """Forward with Post-LN applied to the linear output.""" | ||
| output, bias = super().forward(x) | ||
| return self.post_layernorm(output), bias |
There was a problem hiding this comment.
Post-LayerNorm applied before bias addition may produce incorrect results when bias is present.
TERowParallelLinear.forward returns (output, bias), where bias is a deferred addend applied by the downstream BDA (bias-dropout-add) function — the output tensor does not yet have bias added. The current implementation applies self.post_layernorm(output) before bias is added to it. For EXAONE (where add_bias_linear=False and bias is always None) this is harmless, but any other model or configuration that reuses TERowParallelLinearLayerNorm with add_bias_linear=True will normalize a biasless output and then add bias outside the norm, which is incorrect.
Add a guard or document the invariant explicitly.
🛡️ Proposed defensive fix
def forward(self, x):
"""Forward with Post-LN applied to the linear output."""
output, bias = super().forward(x)
+ if bias is not None:
+ raise ValueError(
+ "TERowParallelLinearLayerNorm does not support deferred bias "
+ "(add_bias_linear must be False for Post-LN architectures)."
+ )
return self.post_layernorm(output), bias🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py` around lines 91 - 94,
TERowParallelLinearLayerNorm.forward currently calls super().forward(x) ->
(output, bias) then applies self.post_layernorm(output) before bias is added;
add a guard to prevent silently incorrect behavior by checking the returned
bias: either assert bias is None with a clear message (e.g., raise
AssertionError("TERowParallelLinearLayerNorm expects no deferred bias; set
add_bias_linear=False or handle bias before layernorm")), or if you want to
support bias, add bias to output before calling self.post_layernorm (i.e.,
output = output + bias) and return (normalized_output, None); make the change
inside TERowParallelLinearLayerNorm.forward and reference super().forward,
self.post_layernorm, and the bias variable.
| # ============================================================================= | ||
|
|
||
|
|
||
| def exaone4_layer_spec(config: "GPTModelProvider") -> ModuleSpec: |
There was a problem hiding this comment.
Unused function argument config in exaone4_layer_spec (confirmed by static analysis).
The function builds a fully static ModuleSpec without referencing config at any point. If it is a forward-looking placeholder for config-dependent specs (e.g., 32B hybrid attention), document that intent with a comment; otherwise remove the parameter.
♻️ Proposed fix (if config is intentionally unused for now)
-def exaone4_layer_spec(config: "GPTModelProvider") -> ModuleSpec:
+def exaone4_layer_spec(_config: "GPTModelProvider") -> ModuleSpec:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def exaone4_layer_spec(config: "GPTModelProvider") -> ModuleSpec: | |
| def exaone4_layer_spec(_config: "GPTModelProvider") -> ModuleSpec: |
🧰 Tools
🪛 Ruff (0.15.2)
[warning] 102-102: Unused function argument: config
(ARG001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py` at line 102, The
function exaone4_layer_spec currently declares a parameter config but never uses
it; either remove the unused parameter or clearly mark it as an intentional
placeholder: if you intend future config-dependent specs, retain the parameter
and add a comment above exaone4_layer_spec explaining that config is
intentionally unused for now (e.g., "placeholder for future config-dependent
layer specs such as 32B hybrid attention"), otherwise change the signature to
def exaone4_layer_spec() -> ModuleSpec and update all call sites that pass a
config to call the new no-arg function.
| params={"attn_mask_type": AttnMaskType.causal}, | ||
| submodules=SelfAttentionSubmodules( | ||
| linear_qkv=TEColumnParallelLinear, # No Pre-LN (pure Post-LN arch) | ||
| core_attention=None, # Use default DotProductAttention |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, find the file mentioned in the review
find . -name "exaone4_provider.py" -type fRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 126
🏁 Script executed:
# Search for SelfAttentionSubmodules definition
rg -n "class SelfAttentionSubmodules" --type py -A 10Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 53
🏁 Script executed:
# Search for core_attention parameter usage and handling
rg -n "core_attention" --type py -B 3 -A 5 | head -100Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 8878
🏁 Script executed:
# Check how SelfAttention uses the core_attention parameter
rg -n "class SelfAttention" --type py -A 20Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 53
🏁 Script executed:
# View the context around line 126 in exaone4_provider.py
sed -n '110,150p' ./src/megatron/bridge/models/exaone/exaone4_provider.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1588
🏁 Script executed:
# Check imports and structure of exaone4_provider.py
head -50 ./src/megatron/bridge/models/exaone/exaone4_provider.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1966
🏁 Script executed:
# Search for how SelfAttentionSubmodules is imported or used
rg -n "SelfAttentionSubmodules" --type pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1550
🏁 Script executed:
# Check for None handling in core_attention across the codebase
rg -n "core_attention.*=.*None" --type pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 204
🏁 Script executed:
# Check gemma2_provider.py usage of SelfAttentionSubmodules
sed -n '310,330p' ./src/megatron/bridge/models/gemma/gemma2_provider.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 982
🏁 Script executed:
# Check gemma3_provider.py usage of SelfAttentionSubmodules
sed -n '210,230p' ./src/megatron/bridge/models/gemma/gemma3_provider.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1032
🏁 Script executed:
# Search for imports in exaone4_provider to see what attention classes are available
grep -n "import\|from" ./src/megatron/bridge/models/exaone/exaone4_provider.py | head -60Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 741
🏁 Script executed:
# Check if there's a DotProductAttention or default import in exaone4_provider.py
grep -n "DotProductAttention\|CoreAttention\|MCoreDot" ./src/megatron/bridge/models/exaone/exaone4_provider.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 150
🏁 Script executed:
# Search all files for core_attention assignments to understand the pattern
rg -n "core_attention=" --type py -B 2 -A 1Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 4302
🏁 Script executed:
# Check if there are any tests that verify core_attention=None behavior
rg -n "core_attention.*None" --type py -B 3 -A 3Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 878
🏁 Script executed:
# Look for DotProductAttention or related classes in the codebase
rg -n "DotProductAttention|TEDotProductAttention" --type pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 4100
Import and explicitly assign a DotProductAttention class to core_attention instead of relying on None handling.
All other models in the codebase (Gemma2, Gemma3, Ministral3, GPT) explicitly assign an attention class to core_attention (e.g., Gemma2DotProductAttention, Gemma3TEDotProductAttention, MCoreDotProductAttention). Exaone4 is the only model using core_attention=None, which assumes megatron-core's SelfAttention handles None by falling back to a default—a behavior that should be verified or explicitly implemented. Import DotProductAttention from megatron.core.transformer.dot_product_attention (like gpt_provider.py does) and assign it to core_attention to match the established pattern.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py` at line 126, Import
DotProductAttention from megatron.core.transformer.dot_product_attention and
explicitly assign it to the core_attention argument instead of leaving
core_attention=None in ExaOne4 provider; update the import list in
src/megatron/bridge/models/exaone/exaone4_provider.py (matching the pattern used
in gpt_provider.py) and set core_attention=DotProductAttention where the ExaOne4
model is constructed so the provider uses the same explicit attention class as
Gemma2/Gemma3/MCore-based providers.
| rotary_percent: float = 1.0 | ||
|
|
||
| # Custom layer spec for Post-LN architecture | ||
| transformer_layer_spec: Union[ModuleSpec, Callable[["GPTModelProvider"], ModuleSpec]] = exaone4_layer_spec |
There was a problem hiding this comment.
Use X | Y union syntax instead of Union[X, Y] (Python 3.10+ requirement).
♻️ Proposed fix
- transformer_layer_spec: Union[ModuleSpec, Callable[["GPTModelProvider"], ModuleSpec]] = exaone4_layer_spec
+ transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = exaone4_layer_specAs per coding guidelines: "Use X | Y for union types instead of Union[X, Y]" and "Conform code to Python 3.10+."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| transformer_layer_spec: Union[ModuleSpec, Callable[["GPTModelProvider"], ModuleSpec]] = exaone4_layer_spec | |
| transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = exaone4_layer_spec |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/models/exaone/exaone4_provider.py` at line 175, Change
the type annotation for transformer_layer_spec to use PEP 604 union syntax:
replace Union[ModuleSpec, Callable[["GPTModelProvider"], ModuleSpec]] with
ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] (symbol:
transformer_layer_spec, types: ModuleSpec, Callable, GPTModelProvider,
exaone4_layer_spec); also remove any now-unused Union import from typing.
Signed-off-by: 김재우 <pewpewplay315@gmail.com>
|
Hi @yaoyu-33 — I hope you're doing well! I wanted to send a gentle ping on this PR. It adds HuggingFace ↔ Megatron checkpoint conversion support for EXAONE 4.0 (LG AI Research), following the existing bridge pattern. The CI checks are currently waiting for workflow approval from a vetter. Would you be able to take a look and trigger the CI when you get a chance? I'd greatly appreciate any feedback. Thank you! |
|
|
||
|
|
||
| @dataclass | ||
| class Exaone4ModelProvider(GPTModelProvider): |
There was a problem hiding this comment.
is there an existing model where these values are different from this default? If not we should write these as default values in the model building blocks. Ideally we can remove the model specific provider if not necessary
|
Sorry for the delay. Can you also add a functional test and unit test following this example https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2602/changes |
|
|
||
|
|
||
| class TERowParallelLinearLayerNorm(TERowParallelLinear): | ||
| """Row-parallel linear with an additional Post-LayerNorm on the output. |
There was a problem hiding this comment.
TERowParallelLinearLayerNorm is already provided in src/megatron/bridge/models/gemma/gemma2_provider.py. could you please take a look to see whether it's different from EXAONE 4.0? if they are similar, it could be pulled out to a common modules file and make both models import the class from there
svcnvidia-nemo-ci
added
the
waiting-on-customer
yaoyu-33
added
area:model
thank you for taking the time to review this PR , i'll add the README/scripts and tests. for |
|
I checked the existing |
|
(additional) i will add the README page, run scripts, and tests , working on it now. |
Move duplicated TERowParallelLinearLayerNorm class into models/common/te_layers.py and update Gemma2, Gemma3, and EXAONE imports. No behavior change on the normal no-bias path; adds a defensive assertion for deferred bias. Signed-off-by: Bias92 <pewpewplay315@gmail.com>
Signed-off-by: Bias92 <pewpewplay315@gmail.com>
Signed-off-by: Bias92 <pewpewplay315@gmail.com>
Signed-off-by: Bias92 <pewpewplay315@gmail.com>
Drop two fields from Exaone4ModelProvider that duplicate parent defaults: - share_embeddings_and_output_weights (parent: True) - rotary_percent (parent: 1.0) Per reviewer feedback on PR NVIDIA-NeMo#2532. Signed-off-by: Bias92 <pewpewplay315@gmail.com>
|
Thank you for the quick update. By the way, We have just released https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/skills/adding-model-support in the repo. Feel free to check out SKILL.md, llm-patterns.md, and tests-and-examples.md for the expected patterns and checklist using any AI tool. hope that also help model support work in the future. |
|
Hi @Bias92 , let me know if it's ready for review. thank you for the input |
|
oh sorry, i was busy for while. i will take a look , today. |
|
no hurry. Just ping me when it's ready for review. |
…ling test The llama3-style rope scaling fields (low_freq_factor, high_freq_factor, original_max_position_embeddings) were dataclass fields on the removed Exaone4ModelProvider but are not declared on GPTModelProvider, so passing them to the constructor raises TypeError and fails Launch_Unit_Tests_Core. Set them as plain attributes instead, mirroring provider_bridge. Signed-off-by: Bias92 <pewpewplay315@gmail.com>
…bias guard - Remove the duplicated "Exaone4Bridge" entry from models/__init__.py __all__ (the alphabetically ordered entry is kept) - Add the NVIDIA Apache 2.0 header to models/exaone/__init__.py - Replace assert with ValueError in TERowParallelLinearLayerNorm.forward so the bias guard survives python -O Signed-off-by: Bias92 <pewpewplay315@gmail.com>
|
oh sorry i'm just back. thanks for the progress |
|
/ok to test 8177979 |
Signed-off-by: Bias92 <pewpewplay315@gmail.com>
Covers the deferred-bias ValueError branch in te_layers.py (the only uncovered line in the PR diff) and gives the shared module its own test file. Signed-off-by: Bias92 <pewpewplay315@gmail.com>
|
updated branch + added the missing test. codecov thing was a stale base issue. can you rerun ci? |
|
/ok to test 45afbe5 |
cuichenx
added
ready-to-merge
|
Rerunning codecov as it might running bundling with other PRs |
|
/ok to test 45afbe5 |
|
thx for the merge! |
8 participants
feat: Add EXAONE 4.0 model bridge (LG AI Research)
Summary
Add HuggingFace ↔ Megatron checkpoint conversion support for EXAONE 4.0 (LG AI Research).
Architecture
EXAONE 4.0 uses a pure Post-LayerNorm architecture that differs from standard Pre-LN models:
Key features:
TERowParallelLinearLayerNorm, following Gemma2 pattern)share_embeddings_and_output_weights=True)Files
Tested with
LGAI-EXAONE/EXAONE-4.0-1.2B(30 layers, 2048 hidden)AutoBridge.from_hf_pretrained()→ provider loads correctlyto_megatron_provider()→num_layers=30, hidden_size=2048, qk_layernorm=Truemapping_registry()→ 10 mappings verifiedFuture work
- EXAONE 4.0 32B: Hybrid attention (LLLG pattern — 3 local + 1 global layers) with sliding window
- K-EXAONE 236B MoE: Expert parallel + MoE bridge support
- Training recipes
feat: Add EXAONE 4.0 model bridge (LG AI Research)Summary
Add HuggingFace ↔ Megatron checkpoint conversion support for [EXAONE 4.0](https://huggingface.co/LGAI-EXAONE/EXAONE-4.0-1.2B) (LG AI Research).
Architecture
EXAONE 4.0 uses a pure Post-LayerNorm architecture that differs from standard Pre-LN models:
Key features:
TERowParallelLinearLayerNorm, following Gemma2 pattern)share_embeddings_and_output_weights=True)Files
exaone4_provider.pyTransformerLayerSubmodulesspecexaone4_bridge.py__init__.pymodels/__init__.pyTested with
LGAI-EXAONE/EXAONE-4.0-1.2B(30 layers, 2048 hidden)AutoBridge.from_hf_pretrained()→ provider loads correctlyto_megatron_provider()→num_layers=30, hidden_size=2048, qk_layernorm=Truemapping_registry()→ 10 mappings verifiedFuture work
Summary by CodeRabbit