Add IsVirtualDeivce() and refactor #6726
Conversation
| data_ = data; | ||
| data_.reset(data.get()); |
There was a problem hiding this comment.
what's the difference between these two?
There was a problem hiding this comment.
Better memory management, this is refactoring to reuse the existing data_ to hold the new data.
| @@ -346,7 +346,7 @@ class XLAGraphExecutor : public torch::lazy::LazyGraphExecutor { | |||
| std::vector<size_t> SetBufferDonors(LoweringContext* lowering_ctx); | |||
|
|
|||
| // We don't use upstream Compile to have BuildInputOutputAliases. | |||
| CompilationResult Compile(const std::vector<XLATensorPtr>& tensors, | |||
| CompilationResult Compile(std::vector<XLATensorPtr>& tensors, | |||
There was a problem hiding this comment.
I can see how this will become problematic for the dynamo, since in dynamo we will first dry run the compilation and don't execute the compiled program. Please leave a TODO somewhere for the dynamo integration with auto-sharding.
There was a problem hiding this comment.
Makes sense I will also add a section in the design doc.
There was a problem hiding this comment.
// We don't use upstream Compile to have BuildInputOutputAliases.
// TODO(yeounoh) auto-sharding can change tensors shardings, which needs to be
// accounted for in Dynamo integration.
CompilationResult Compile(std::vector<XLATensorPtr>& tensors,
absl::Span<const std::string> devices,
const SyncTensorCollection& coll,
PostOrderData* po_data,
const std::vector<torch::lazy::Value>& ir_values);
| sharding_specs.push_back(""); | ||
| } | ||
| sharding_specs.push_back( | ||
| GetXLAShardingSpec(bridge::GetXlaTensor(tensor))); |
There was a problem hiding this comment.
what will GetXLAShardingSpec return when sharding_spec is null?
There was a problem hiding this comment.
"" empty string as expected.
yeounoh
force-pushed
the
add_virtual_deivce_utils
branch
from
11cc726 to
381077a
Compare
yeounoh
force-pushed
the
add_virtual_deivce_utils
branch
from
381077a to
b21f260
Compare
This addresses the following to support #6322
tensorsargument toXLAGraphExecutor::Compilemethod, since they can be updated during auto-sharding paass.