Broaden packages.find.include to expose megatron.training#33
Closed
Shi-Dong wants to merge 3 commits into
Closed
Conversation
In addition to `megatron.core`, also expose `megatron.training` (and its sub-packages) via setuptools' package discovery, so that `pip install -e .` against this checkout makes `from megatron.training import ...` work without callers having to set PYTHONPATH. Why: miles imports from `megatron.training.arguments`, `megatron.training.checkpointing`, `megatron.training.global_vars`, `megatron.training.tokenizer.tokenizer`, and `megatron.training.training`. Until now those imports only worked because the production environment adds `/root/Megatron-LM` (the editable install's root) to sys.path, which is a side effect of `pip install -e` rather than a deliberate package export. miles' CPU CI runner makes that side effect explicit via `PYTHONPATH=$workspace:$workspace/third_party/Megatron-LM`, but that override doesn't survive packaging miles as a wheel for the eventual `pip install miles` goal. By exposing `megatron.training.*` in pyproject.toml, the side-effect dependency goes away. Callers can rely on the explicit Python package list instead of the editable-install path trick. This change does not affect runtime semantics for code that was already importing from `megatron.training` (those callers continue to find the same modules); it only changes which sys.path entries make the imports resolve, broadening it from "editable-install root only" to "any pip install location (editable, wheel, or otherwise)".
3 tasks
added 2 commits
May 24, 2026 09:04
megatron.training.initialize unconditionally imports from megatron.legacy at module-load time, so any narrower include list (e.g. exposing only core+training) breaks transitively. Expose all megatron.* sub-packages instead.
megatron, megatron.legacy, and a few sub-packages have no __init__.py — they are namespace packages. The default find_packages mode skips namespace packages, so the previous broaden-to-megatron* commit did not actually expose megatron.legacy. Enable namespaces=true so the discovery is namespace-package-aware. Now pip install -e exposes the full megatron.* namespace without relying on the editable-install side effect that adds the project root to sys.path.
2 tasks
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
[tool.setuptools.packages.find].includeto also exposemegatron.trainingandmegatron.training.*(in addition to the existingmegatron.core/megatron.core.*).Why
miles imports from several modules under
megatron.training.*(arguments,checkpointing,global_vars,tokenizer.tokenizer,training). Until now those imports succeed only becausepip install -e .adds the project root tosys.pathas a side effect — not becausemegatron.trainingis an explicitly declared package.miles' CPU CI runner makes that side effect explicit via
PYTHONPATH=$workspace/third_party/Megatron-LM. That PYTHONPATH workaround doesn't survive packaging miles as a wheel (the eventualpip install milesgoal in the broader Phase 2 roadmap).Declaring
megatron.training.*in the explicit package list lets a normalpip install(editable, wheel, or otherwise) expose it. No code in this PR changes; only the discovery list does.What this does NOT change
megatron.training: the same modules resolve.megatron-core's project name. The wheel is stillmegatron-core, just with a slightly broader set of sub-packages.Test plan
third_party/Megatron-LMsubmodule pointer to this branch's HEAD and removes thePYTHONPATH=$workspace/third_party/Megatron-LMoverride from_run-ci.yml. CI on that PR is the validation gate: stage-a-fast (CPU runner) must pass — that's the path that previously needed the PYTHONPATH hack to findmegatron.training.unset PYTHONPATH && python -c "import megatron.training; import megatron.training.training"succeeds.Stacked on
radixark/Megatron-LM PR #32 (Phase 1: remove
miles_megatron_plugins/). After both this PR and #32 merge, themiles-mainbranch carries both changes; the miles repo's submodule pointer flips back tobranch = miles-main.