Skip to content
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.

Sign up for GitHub

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

fix xla with multiple devices #2081

Merged
merged 2 commits into from
Oct 28, 2024
Merged

fix xla with multiple devices #2081

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5137816
add device assignment
borg323 Oct 22, 2024
c7bead2
don't set device_ordinal
borg323 Oct 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions src/neural/xla/hlo.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,15 @@ message CompileEnvOptionProto {
required OptionOverrideProto value = 2;
}

message XlaDeviceAssignmentProto {
optional int32 replica_count = 1;
optional int32 computation_count = 2;
message ComputationDevice {
repeated int64 replica_device_ids = 1;
}
repeated ComputationDevice computation_devices = 3;
}

message ExecutableBuildOptionsProto {
// If set, this is the device to build the computation for. Valid
// device_ordinal values are: 0 to # of devices - 1. These values are
Expand Down Expand Up @@ -386,6 +395,12 @@ message ExecutableBuildOptionsProto {
// Whether HLOs should be deduplicated.
optional bool deduplicate_hlo = 8;

// If set, this specifies a static device assignment for the computation.
// Otherwise, the computation will be compiled generically and can be run with
// any device assignment compatible with the computation's replica and
// partition counts.
optional XlaDeviceAssignmentProto device_assignment = 9;

// Whether input and output buffers are aliased if the associated parameter is
// passed-through XLA modules without being changed.
optional bool alias_passthrough_params = 10;
Expand Down
11 changes: 10 additions & 1 deletion src/neural/xla/xla_runner.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,16 @@ void XlaRunner::AddModule(size_t minibatch_size,
pblczero::CompileOptionsProto options;
options.mutable_executable_build_options()->set_num_replicas(1);
options.mutable_executable_build_options()->set_num_partitions(1);
options.mutable_executable_build_options()->set_device_ordinal(device_);
options.mutable_executable_build_options()
->mutable_device_assignment()
->set_replica_count(1);
options.mutable_executable_build_options()
->mutable_device_assignment()
->set_computation_count(1);
options.mutable_executable_build_options()
->mutable_device_assignment()
->add_computation_devices()
->add_replica_device_ids(device_);
auto executable = pjrt_client_->CompileHlo(module.OutputAsString(),
options.OutputAsString());
executables_.push_back({minibatch_size, std::move(executable)});
Expand Down