Skip to content

Add a helper class to handle mesh and sharding #8967

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 16, 2025

Conversation

qihqi
Copy link
Collaborator

@qihqi qihqi commented Apr 12, 2025

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
Copy link
Collaborator

@tengyifei tengyifei left a comment

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

@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
24 checks passed
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