Skip to content

[Context Parallel] Add support for custom device mesh#13064

Open
DN6 wants to merge 1 commit intomainfrom
custom-device-map
Open

[Context Parallel] Add support for custom device mesh#13064
DN6 wants to merge 1 commit intomainfrom
custom-device-map

Conversation

@DN6
Copy link
Collaborator

@DN6 DN6 commented Feb 2, 2026

What does this PR do?

Allow users to set a custom device mesh to use with CP.

  mesh = init_device_mesh("cuda", mesh_shape=(2, 1, 4), mesh_dim_names=("ring", "ulysses", "fsdp"))
  cp_config = ContextParallelConfig(ring_degree=2, mesh=mesh)
  pipe.transformer.enable_parallelism(config=cp_config)

Fixes # (issue)
#12880

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@DN6
Copy link
Collaborator Author

DN6 commented Feb 2, 2026

cc: @pthombre would you be able to test to see if this works for your use case?

@DN6 DN6 added the roadmap Add to current release roadmap label Feb 2, 2026
@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@DefTruth
Copy link
Contributor

DefTruth commented Feb 5, 2026

cool, I am also looking for this feature~

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

Labels

roadmap Add to current release roadmap

Projects

Status: In Progress

Development

Successfully merging this pull request may close these issues.

3 participants