Skip to content

Add a helper class to handle mesh and sharding - #8967

Merged
qihqi merged 3 commits into
masterfrom
hanq_xla2_model_constructor
Apr 16, 2025
Merged

Add a helper class to handle mesh and sharding#8967
qihqi merged 3 commits into
masterfrom
hanq_xla2_model_constructor

Conversation

@qihqi

@qihqi qihqi commented Apr 12, 2025

Copy link
Copy Markdown
Collaborator

The first functionality is to initialize a model
on sharded devices

fixes #8635

@qihqi
qihqi requested a review from tengyifei April 12, 2025 00:43

@tengyifei tengyifei left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

lgtm but need to address the comments

Comment thread torchax/torchax/mesh_util.py Outdated
Comment thread torchax/torchax/mesh_util.py Outdated
Comment thread torchax/torchax/mesh_util.py
Comment thread torchax/test_dist/test_mesh_util.py
@qihqi
qihqi force-pushed the hanq_xla2_model_constructor branch 2 times, most recently from 9edc779 to 7fef690 Compare April 14, 2025 19:17
@qihqi
qihqi enabled auto-merge (squash) April 14, 2025 23:14
@qihqi
qihqi force-pushed the hanq_xla2_model_constructor branch from 7fef690 to d41bf43 Compare April 15, 2025 19:04
qihqi added 3 commits April 16, 2025 10:56
The first functionality is to initialize a model
on sharded devices
@qihqi
qihqi force-pushed the hanq_xla2_model_constructor branch from d41bf43 to 76f86d6 Compare April 16, 2025 17:56
@qihqi
qihqi merged commit 7882475 into master Apr 16, 2025
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.

[torchax] jit compile the model constructor

2 participants