Skip to content

merge to main of feat(flux): core runtime + Megatron adapter scaffolding#856

Merged
Xiaoming-AMD merged 13 commits into
mainfrom
feat/flux/core
Jul 7, 2026
Merged

merge to main of feat(flux): core runtime + Megatron adapter scaffolding#856
Xiaoming-AMD merged 13 commits into
mainfrom
feat/flux/core

Conversation

@luiza-amd

@luiza-amd luiza-amd commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Important: The content of this PR was approved and merged #807, but not to main (auto-target to main was not triggered and the lesson was taken into account for further PRs).
The purpose of this PR is to finalize the merge to main and no new changes were introduced.

Base of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets feat/flux/ci-env and auto-retargets to main once that merges. Every other content PR stacks on this one.

What this changes

The shared runtime scaffolding the rest of the feature builds on: core runtime state + train-runtime wiring, the Megatron adapter and base/pretrain trainers, and a patch auto-loader (patches/__init__.py) that imports every *_patches.py in the package on import, so each later layer only drops in its own patch file with no registry edit. Also carries the shared test root (tests/conftest.py, tests/utils.py) and a one-line .gitignore change (the bare data ignore → root-anchored /data/*, and nothing else) that stops Git from ignoring the in-repo data/ source and config directories the later layers add. Kept deliberately small so it can land first.

Dependencies

Sequenced after the CI-pins PR (feat/flux/ci-env); no functional dependency on the turbo bump.

Test plan

pytest tests/unit_tests/core tests/unit_tests/backends/megatron; lint/pre-commit clean. Validated locally on an AMD GPU container: 58 passed.

Files

39 (core runtime, Megatron adapter/trainers, base patch loader, shared test root, .gitignore).

Flux Split and others added 12 commits June 28, 2026 02:06
… on draft PRs

Bump PRIMUS_TURBO_COMMIT to Primus-Turbo main (0.3.x, validated build) which
exposes the mxfp4 gemm_fp4_impl(..., preshuffled=...) fast path, and
PRIMUS_TURBO_AITER_COMMIT to AITER v0.1.14.post1. Mirror the turbo commit into
benchmark.yaml; the docker build receives both via ci.yaml build-args.

Add a draft-skip guard to the expensive jobs (build-docker, run-unittest-torch,
run-unittest-jax) so the stacked Flux PRs can sit in Draft without firing the
full pipeline; code-lint stays ungated. Add ready_for_review to the
pull_request trigger types so flipping a draft to Ready triggers the first run.

MEGATRON_PATH export intentionally omitted: the Flux diffusion unit tests run
in-process (no megatron-importing subprocess remains after the te_vs_local
in-process conversion), and tests/conftest.py auto-adds
third_party/Megatron-LM to sys.path when MEGATRON_PATH is unset. CI already
checks out submodules: recursive, so megatron imports without any extra env.
… DeepEP pad_multiple

The branch's pins pointed at older commits missing required APIs, breaking
DeepEP MoE training with:
    DeepEPTokenDispatcher.__init__() got an unexpected keyword argument 'pad_multiple'

- PRIMUS_TURBO_COMMIT 3cd482de -> bc6c7d2a :
  old pin lacked both gemm_fp4_impl(preshuffled=) and
  DeepEPTokenDispatcher(pad_multiple=); new pin has both.
- PRIMUS_TURBO_AITER_COMMIT f299f579 -> 0f3c58e6: the real v0.1.14.post1
  tag commit required by Primus-Turbo main's aiter_utils.py.
… on draft PRs

Bump PRIMUS_TURBO_COMMIT to Primus-Turbo main (0.3.x, validated build) which
exposes the mxfp4 gemm_fp4_impl(..., preshuffled=...) fast path, and
PRIMUS_TURBO_AITER_COMMIT to AITER v0.1.14.post1. Mirror the turbo commit into
benchmark.yaml; the docker build receives both via ci.yaml build-args.

Add a draft-skip guard to the expensive jobs (build-docker, run-unittest-torch,
run-unittest-jax) so the stacked Flux PRs can sit in Draft without firing the
full pipeline; code-lint stays ungated. Add ready_for_review to the
pull_request trigger types so flipping a draft to Ready triggers the first run.

MEGATRON_PATH export intentionally omitted: the Flux diffusion unit tests run
in-process (no megatron-importing subprocess remains after the te_vs_local
in-process conversion), and tests/conftest.py auto-adds
third_party/Megatron-LM to sys.path when MEGATRON_PATH is unset. CI already
checks out submodules: recursive, so megatron imports without any extra env.

rebase to main branch
@luiza-amd luiza-amd self-assigned this Jul 7, 2026
@luiza-amd luiza-amd marked this pull request as draft July 7, 2026 06:58
Comment thread primus/backends/megatron/__init__.py Dismissed
Comment thread primus/backends/megatron/megatron_base_trainer.py Dismissed
Comment thread primus/backends/megatron/megatron_base_trainer.py Dismissed
Comment thread primus/core/runtime/train_runtime.py Dismissed
@luiza-amd luiza-amd marked this pull request as ready for review July 7, 2026 07:35
@Xiaoming-AMD Xiaoming-AMD merged commit e2f7ee6 into main Jul 7, 2026
12 of 13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants