Skip to content

feat(flux): diffusion data preprocessing pipelines + data CLI#817

Draft
jasainio wants to merge 6 commits into
feat/flux/datafrom
feat/flux/prep
Draft

feat(flux): diffusion data preprocessing pipelines + data CLI#817
jasainio wants to merge 6 commits into
feat/flux/datafrom
feat/flux/prep

Conversation

@jasainio

Copy link
Copy Markdown
Contributor

Part 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/data — review after it.

What this changes

The offline data-preprocessing layer (auth / download / finalize / validate plus the raw/ingest/encoded pipeline stages), wired into a new primus data CLI subcommand.

Dependencies

Sequenced after the CI-pins PR (feat/flux/ci-env); builds on feat/flux/data (uses its providers and inherits the energon/webdataset deps).

Test plan

pytest tests/unit_tests/backends/megatron/diffusion/data/preprocessing tests/unit_tests/cli/test_data_config.py. Validated locally on an AMD GPU container: 63 passed.

Files

20 (preprocessing pipelines, data CLI subcommand + tests).

Flux Split and others added 6 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.
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.

1 participant