Skip to content

feat(flux): common diffusion module (embeddings, normalization, DiT block)#810

Draft
jasainio wants to merge 1 commit into
feat/flux/corefrom
feat/flux/model-common
Draft

feat(flux): common diffusion module (embeddings, normalization, DiT block)#810
jasainio wants to merge 1 commit into
feat/flux/corefrom
feat/flux/model-common

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/core — review after it. Parent of the Flux model PR.

What this changes

The model-agnostic diffusion building blocks shared by all diffusion models: the common diffusion module, diffusion config/embeddings/normalization, and the diffusion transformer (DiT) block. Also lands the shared diffusion test scaffolding (conftest/helpers/constants) the later diffusion PRs reuse.

Dependencies

Sequenced after the CI-pins PR (feat/flux/ci-env); builds on feat/flux/core.

Test plan

pytest tests/unit_tests/backends/megatron/diffusion -k "embeddings or normalization". Validated locally on an AMD GPU container: 7 passed.

Files

14 (common diffusion module, embeddings/normalization/config, DiT block, shared diffusion test scaffolding).

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