Skip to content

feat: support the wan22's dit and tp parallel.#1445

Open
ethan686 wants to merge 3 commits into
jd-opensource:mainfrom
ethan686:wan22_push_dit_tp
Open

feat: support the wan22's dit and tp parallel.#1445
ethan686 wants to merge 3 commits into
jd-opensource:mainfrom
ethan686:wan22_push_dit_tp

Conversation

@ethan686
Copy link
Copy Markdown
Contributor

including three changes, dit transformer impl, tp parallel for dit model and all the model args changes for wan2.2. Note: use mindiesd's tensor as dit's input, no tp, 10 transformer layer, the cos similarity is 100% match with mindiesd.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request implements the Wan 2.2 DiT model, introducing the transformer architecture and expanding model arguments to support VAE and DiT parameters. It also refactors DiTParallelLinear to include Megatron-style tensor parallelism. Reviewer feedback highlights multiple style guide violations, such as the requirement for explicit constructors, final class markers, fixed-width integers, and specific naming conventions for enums and parameters. Additionally, the reviewer identified a redundant semicolon and noted that structs containing member functions must be converted to classes.

Comment thread xllm/models/dit/npu/wan2_2/transformer_wan2_2.h Outdated
Comment thread xllm/models/dit/npu/wan2_2/transformer_wan2_2.h Outdated
Comment thread xllm/models/dit/npu/wan2_2/transformer_wan2_2.h Outdated
Comment thread xllm/models/dit/npu/wan2_2/transformer_wan2_2.h Outdated
Comment thread xllm/models/dit/npu/wan2_2/transformer_wan2_2.h Outdated
Comment thread xllm/models/dit/utils/dit_parallel_linear.h
Comment thread xllm/models/dit/utils/dit_parallel_linear.h
Comment thread xllm/models/dit/utils/common_util.h Outdated
Copy link
Copy Markdown
Collaborator

@yq33victor yq33victor left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

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.

3 participants