Skip to content

Parallarize some of the SPMD logic - #5432

Closed
JackCaoG wants to merge 4 commits into
masterfrom
JackCaoG/dynamo_spmd_optimization
Closed

JackCaoG wants to merge 4 commits into
masterfrom
JackCaoG/dynamo_spmd_optimization

Conversation

@JackCaoG

Copy link
Copy Markdown
Collaborator

With SPMD one process needs to handle all device within a single host, this pr tries to parallelize some of the per-device operation.

@JackCaoG
JackCaoG requested a review from yeounoh August 10, 2023 01:01
@JackCaoG
JackCaoG marked this pull request as ready for review August 10, 2023 17:50
@JackCaoG
JackCaoG requested a review from will-cromar August 10, 2023 17:50
Comment thread torch_xla/csrc/xla_sharding_util.cpp Resolved
Comment thread torch_xla/csrc/init_python_bindings.cpp Resolved
coll.device = device;
coll.unlocker = DeviceLockerArena::Get()->LockDevices({device});
// init this before the lock to save some time.
// std::vector<torch::lazy::BackendDataPtr> arguments(graph_inputs.size());

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.

nit. Remove this commented code?

Comment thread torch_xla/csrc/xla_graph_executor.cpp Resolved
// GetXlaData must be called within a lock region, otherwise it might
// extract the placeholder inserted by previous execution.
TORCH_LAZY_TIMED("RunCachedGraphInputData");
// size_t per_thread_work = 200;

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 see a bunch of commented lines here -- maybe we should clean them before mering?

Comment thread torch_xla/csrc/xla_sharding_util.h Resolved
Comment thread torch_xla/csrc/xla_sharding_util.cpp Resolved
Comment thread torch_xla/csrc/xla_sharding_util.h Resolved

@yeounoh yeounoh left a comment

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.

Looks good, thanks @JackCaoG . Left some comments.

@JackCaoG

Copy link
Copy Markdown
Collaborator Author

@yeounoh This pr has been too big, let me split it up and open smaller ones.

@JackCaoG

Copy link
Copy Markdown
Collaborator Author

I break this pr into #5446, #5447 and 13f9382. Closing this one.

@JackCaoG JackCaoG closed this Aug 15, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

3 participants