Skip to content

fix async tp - #4045

Open
acisseJZhong wants to merge 2 commits into
gh/acisseJZhong/26/basefrom
gh/acisseJZhong/26/head
Open

fix async tp#4045
acisseJZhong wants to merge 2 commits into
gh/acisseJZhong/26/basefrom
gh/acisseJZhong/26/head

Conversation

@acisseJZhong

@acisseJZhong acisseJZhong commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Summary

inheriting from #3869
Fix async TP setup across SPMD backends by:

  • enabling symmetric memory on the exact dense TP process group
  • adding ParallelDims.get_dense_tp_mesh() for backend-aware mesh
    selection
  • updating all async TP callsites
  • propagating async TP and sequence parallel settings through
    inference config

[ghstack-poisoned]
[ghstack-poisoned]
acisseJZhong added a commit that referenced this pull request Aug 1, 2026
ghstack-source-id: e47ce56
Pull-Request: #4045
)

with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

what does this do?

context_parallel_degree=1,
pipeline_parallel_degree=1,
enable_sequence_parallel=False,
enable_sequence_parallel=self.enable_sequence_parallel,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

we are not convinced by evidence that this is needed at all, let's wait

Comment on lines +82 to +84
maybe_enable_async_tp(
parallelism, compile_config, parallel_dims.get_dense_tp_mesh()
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I feel we can put this under compile enablement -- if model compile is enabled and inductor backend is used, we should always enable async TP (any reason not to?). And we should delete the enable_async_tp flag. WDYT?

My eventual goal is to delete the special parallelize_deepseekv3, having a unified one for decoder, and make it simple.

spmd_backend: Literal["default", "spmd_types"] = "default"
"""SPMD backend used by TorchTitan model parallelization in the generator."""

enable_async_tensor_parallel: bool = False

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

similar

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/rl ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants