Fix NameError in pretrain_retro.py (add import_module), remove unused… #2084
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
I fixed a runtime error and cleaned up minor issues in pretrain_retro.py:
What I changed
Added missing import:
from importlib import import_module Reason: import_module(...) is called in core_model_provider when args.spec is set; without this import the code throws a NameError at runtime.
Removed unused import:
MegatronTokenizer (unused) — removed to avoid noise and potential lint warnings.
Improved log message:
Changed print_rank_0('building GPT model ...') → print_rank_0('building Retro model ...') to accurately reflect the model being constructed.
Why
The missing import is a definite runtime bug that will break the experimental spec-loading path (args.spec).
Removing unused imports helps keep the codebase clean and avoids linter complaints.
Updating the log message reduces confusion when reading logs (this script builds a Retro model, not GPT).
Testing / How I validated
Static review to confirm import_module usage and the missing import.
Confirmed MegatronTokenizer is not referenced elsewhere in the file.
Confirmed the corrected log message matches the behavior of core_model_provider.
Risk / Backward compatibility
Low risk: only adds an import, removes an unused import, and updates a log string. No behavior/functionality changes to model logic.
If other modules relied on MegatronTokenizer being imported at this module level (unlikely), we should import it explicitly from megatron.core.tokenizers.