-
Notifications
You must be signed in to change notification settings - Fork 5.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[AutoParallel] Adapt static spmd rules for dynamic graph (#56367)
* move matmul spmd rules into phi * add basic infer spmd utils * addspmd factory * fix compile error * add unittest * refine infer spmd test and utils * debug infer spmd test * adapt python test * poish details * change to vector attr arg * revert needless change * update matmul spmd rule test * remove original rule * polish details * fix marco error * add comment * pass backward test * fix compile error * add cmake rule for spmd_rules_test * add dist meta tensor * update pybind impl * add marco for rules
- Loading branch information
Showing
21 changed files
with
1,022 additions
and
189 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
51 changes: 51 additions & 0 deletions
51
paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.cc
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. | ||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. */ | ||
|
|
||
| #include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h" | ||
|
|
||
| #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h" | ||
|
|
||
| namespace phi { | ||
| namespace distributed { | ||
|
|
||
| phi::DDim DistMetaTensor::dims() const { | ||
| // member values in tensor_ have higher priority than those in DistMetaTensor | ||
| if (tensor_ != nullptr) { | ||
| PADDLE_ENFORCE_EQ(this->is_dist(), | ||
| true, | ||
| phi::errors::InvalidArgument( | ||
| "The current MetaTensor doesn't contains " | ||
| "DistTensor when call `dist_attr` method.")); | ||
| return MetaTensor::dims(); | ||
| } else { | ||
| return dims_; | ||
| } | ||
| } | ||
|
|
||
| const distributed::TensorDistAttr& DistMetaTensor::dist_attr() const { | ||
| // member values in tensor_ have higher priority than those in DistMetaTensor | ||
| if (tensor_ != nullptr) { | ||
| PADDLE_ENFORCE_EQ(this->is_dist(), | ||
| true, | ||
| phi::errors::InvalidArgument( | ||
| "The current MetaTensor doesn't contains " | ||
| "DistTensor when call `dist_attr` method.")); | ||
| return static_cast<phi::distributed::DistTensor*>(tensor_)->dist_attr(); | ||
| } else { | ||
| return dist_attr_; | ||
| } | ||
| } | ||
|
|
||
| } // namespace distributed | ||
| } // namespace phi |
68 changes: 68 additions & 0 deletions
68
paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. | ||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. */ | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h" | ||
| #include "paddle/phi/core/meta_tensor.h" | ||
|
|
||
| namespace phi { | ||
| namespace distributed { | ||
|
|
||
| class DistMetaTensor : public MetaTensor { | ||
| public: | ||
| // supporting implicit construction is easier to use | ||
| DistMetaTensor(TensorBase* tensor) // NOLINT | ||
| : MetaTensor(tensor) {} | ||
| DistMetaTensor(const TensorBase& tensor) // NOLINT | ||
| : MetaTensor(tensor) {} | ||
| DistMetaTensor(const TensorBase* tensor) // NOLINT | ||
| : MetaTensor(tensor) {} | ||
| DistMetaTensor(TensorBase& tensor) // NOLINT | ||
| : MetaTensor(tensor) {} | ||
| // For static mode only | ||
| DistMetaTensor(const phi::DDim& dims, const TensorDistAttr& dist_attr) | ||
| : dims_(dims), dist_attr_(dist_attr) {} | ||
|
|
||
| DistMetaTensor(DistMetaTensor&&) = default; | ||
| DistMetaTensor& operator=(DistMetaTensor&&) = default; | ||
| DistMetaTensor(const DistMetaTensor&) = default; | ||
| DistMetaTensor& operator=(const DistMetaTensor&) = default; | ||
|
|
||
| virtual ~DistMetaTensor() = default; | ||
|
|
||
| DDim dims() const override; | ||
|
|
||
| const distributed::TensorDistAttr& dist_attr() const; | ||
|
|
||
| private: | ||
| /** | ||
| * Note: When using the semi-automatic parallel segmentation derivation rules | ||
| * of the static graph, in order to facilitate the packaging of the input | ||
| * parameters of the construction, the DistMetaTensor is inherited and | ||
| * encapsulated, and the class members dims_ and dist_attr_ are added to it. | ||
| * | ||
| * The information contained in these two members is also in the tensor of the | ||
| * meta_tensor of the base class, and there is redundancy. | ||
| * | ||
| * We need to pay attention when using it to ensure the consistency. | ||
| * These two members are read-only, and their values cannot be changed | ||
| * after construction. To change their values, they need to be set | ||
| * directly in tensor_*/ | ||
| phi::DDim dims_; | ||
| TensorDistAttr dist_attr_; | ||
| }; | ||
|
|
||
| } // namespace distributed | ||
| } // namespace phi |
97 changes: 97 additions & 0 deletions
97
paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| /* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. | ||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. */ | ||
|
|
||
| #include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h" | ||
|
|
||
| namespace phi { | ||
| namespace distributed { | ||
|
|
||
| void InferSpmdContext::EmplaceBackInput(DistMetaTensor input) { | ||
| inputs_.emplace_back(std::move(input)); | ||
| } | ||
|
|
||
| void InferSpmdContext::EmplaceBackAttr(Attribute attr) { | ||
| attrs_.emplace_back(std::move(attr)); | ||
| } | ||
|
|
||
| const DistMetaTensor& InferSpmdContext::InputAt(size_t idx) const { | ||
| return inputs_.at(idx); | ||
| } | ||
|
|
||
| template <typename AttrType> | ||
| AttrType InferSpmdContext::AttrAt(size_t idx) const { | ||
| try { | ||
| return paddle::get<AttrType>(attrs_.at(idx)); | ||
| } catch (paddle::bad_variant_access const& e) { | ||
| PADDLE_THROW(phi::errors::InvalidArgument( | ||
| "Attribute cast error in InferSpmd Context, the input attr type is " | ||
| "`%s`, but the expected attribute type is `%s`.", | ||
| attrs_.at(idx).type().name(), | ||
| std::type_index(typeid(AttrType)).name())); | ||
| } | ||
| } | ||
|
|
||
| template <> | ||
| bool InferSpmdContext::AttrAt<bool>(size_t idx) const { | ||
| try { | ||
| auto attr = attrs_.at(idx); | ||
| if (attr.type() == typeid(int)) { | ||
| return static_cast<bool>(paddle::get<int>(attr)); | ||
| } else { | ||
| return paddle::get<bool>(attr); | ||
| } | ||
| } catch (paddle::bad_variant_access const& e) { | ||
| PADDLE_THROW(phi::errors::InvalidArgument( | ||
| "Attribute cast error in InferSpmd Context, the input attr type is " | ||
| "`%s`, but the expected attribute type is `bool`.", | ||
| attrs_.at(idx).type().name())); | ||
| } | ||
| } | ||
|
|
||
| const Attribute& InferSpmdContext::AttrAt(size_t idx) const { | ||
| return attrs_.at(idx); | ||
| } | ||
|
|
||
| SpmdRuleFactory& SpmdRuleFactory::Instance() { | ||
| static SpmdRuleFactory g_spmd_rule_map; | ||
| return g_spmd_rule_map; | ||
| } | ||
|
|
||
| bool SpmdRuleFactory::ContainsSpmdRule(const std::string& kernel_name) const { | ||
| return spmd_rule_map_.count(kernel_name) > 0; | ||
| } | ||
|
|
||
| int SpmdRuleFactory::InsertSpmdRule(std::string kernel_name, SpmdRule rule) { | ||
| PADDLE_ENFORCE_NE( | ||
| ContainsSpmdRule(kernel_name), | ||
| true, | ||
| phi::errors::AlreadyExists( | ||
| "`%s` Kernel's Spmd rules has been registered.", kernel_name)); | ||
| spmd_rule_map_.insert({std::move(kernel_name), std::move(rule)}); | ||
| return 0; | ||
| } | ||
|
|
||
| const SpmdRule& SpmdRuleFactory::GetSpmdRule( | ||
| const std::string& kernel_name) const { | ||
| auto it = spmd_rule_map_.find(kernel_name); | ||
| PADDLE_ENFORCE_NE( | ||
| it, | ||
| spmd_rule_map_.end(), | ||
| phi::errors::NotFound("`%s` Kernel's Spmd rules is not registered.", | ||
| kernel_name)); | ||
| return it->second; | ||
| } | ||
|
|
||
| } // namespace distributed | ||
| } // namespace phi |
Oops, something went wrong.