lichenever
4 years ago
parent
a321f402c8
commit
2e1c43483e
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,72 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_CCSRC_FRONTEND_PARALLEL_PIPELINE_TRANSFORMER_PIPELINE_TRANSFORMER_H_
|
||||
#define MINDSPORE_CCSRC_FRONTEND_PARALLEL_PIPELINE_TRANSFORMER_PIPELINE_TRANSFORMER_H_
|
||||
|
||||
#include <utility>
|
||||
#include "ir/value.h"
|
||||
#include "ir/graph_utils.h"
|
||||
#include "base/base.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace parallel {
|
||||
typedef struct {
|
||||
ValueListPtr shape;
|
||||
TypePtr type;
|
||||
AnfNodePtr depend;
|
||||
} SendAttr;
|
||||
|
||||
class PipelineTransformer {
|
||||
public:
|
||||
PipelineTransformer(const FuncGraphManagerPtr &manager, const int &stage, const FuncGraphPtr &root,
|
||||
const int64_t &global_rank, const int64_t &per_stage_rank_num)
|
||||
: manager_(manager),
|
||||
stage_(stage),
|
||||
root_(root),
|
||||
global_rank_(global_rank),
|
||||
per_stage_rank_num_(per_stage_rank_num) {}
|
||||
void Coloring();
|
||||
void BroadCastColoring();
|
||||
void HandleSharedParameter();
|
||||
void CutGraph();
|
||||
void ParameterColoring();
|
||||
void CoverSensShape();
|
||||
void ElimGraphStage();
|
||||
void ElimParameter();
|
||||
|
||||
private:
|
||||
void DoBroadCast(const FuncGraphPtr &func);
|
||||
SendAttr InsertSend(const FuncGraphPtr &graph, const AnfNodePtr ¶meter, const int &user_node_stage,
|
||||
const int &node_stage);
|
||||
void InsertReceive(const FuncGraphPtr &graph, const AnfNodePtr &node, const AnfNodePtr &use_node, const int &index,
|
||||
const int &user_node_stage, const int &node_stage);
|
||||
void CutBorder(const FuncGraphPtr &graph);
|
||||
void ElimRootParameter();
|
||||
bool IsStageNode(const CNodePtr &node);
|
||||
std::pair<CNodePtr, FuncGraphPtr> FindSensNode();
|
||||
FuncGraphManagerPtr manager_;
|
||||
int64_t stage_;
|
||||
FuncGraphPtr root_;
|
||||
int64_t global_rank_;
|
||||
int64_t per_stage_rank_num_;
|
||||
TypePtr type_ptr_;
|
||||
ValueListPtr shape_;
|
||||
};
|
||||
} // namespace parallel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_FRONTEND_PARALLEL_PIPELINE_TRANSFORMER_PIPELINE_TRANSFORMER_H_
|
||||
@ -0,0 +1,99 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* 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 <string>
|
||||
#include <memory>
|
||||
#include "pipeline/jit/pipeline_split.h"
|
||||
#include "utils/ms_context.h"
|
||||
#include "utils/comm_manager.h"
|
||||
#include "frontend/parallel/context.h"
|
||||
#include "frontend/parallel/pipeline_transformer/pipeline_transformer.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace pipeline {
|
||||
|
||||
static int64_t GetRank();
|
||||
static int64_t InferStage(const int64_t &rank_id, const int64_t &stage_num, const int64_t &device_num);
|
||||
static int64_t GetRank() {
|
||||
auto ms_context = MsContext::GetInstance();
|
||||
MS_EXCEPTION_IF_NULL(ms_context);
|
||||
std::string world_group;
|
||||
std::string backend = ms_context->get_param<std::string>(MS_CTX_DEVICE_TARGET);
|
||||
if (backend == kAscendDevice) {
|
||||
world_group = parallel::HCCL_WORLD_GROUP;
|
||||
} else if (backend == kGPUDevice) {
|
||||
world_group = parallel::NCCL_WORLD_GROUP;
|
||||
} else {
|
||||
MS_LOG(EXCEPTION) << "Invalid backend: " << backend;
|
||||
}
|
||||
int64_t global_rank = parallel::ParallelContext::GetInstance()->global_rank();
|
||||
uint32_t rank_id;
|
||||
if (!parallel::ParallelContext::GetInstance()->global_rank_is_set()) {
|
||||
if (!CommManager::GetInstance().GetRankID(world_group, &rank_id)) {
|
||||
MS_LOG(EXCEPTION) << "Get rank id failed.";
|
||||
}
|
||||
global_rank = UintToInt(rank_id);
|
||||
}
|
||||
return global_rank;
|
||||
}
|
||||
|
||||
static int64_t InferStage(const int64_t &rank_id, const int64_t &stage_num, const int64_t &device_num) {
|
||||
if (device_num % stage_num != 0) {
|
||||
MS_LOG(EXCEPTION) << "Device_num must be divisible by the stage_num, got device_num: " << device_num
|
||||
<< "stage_num: " << stage_num;
|
||||
}
|
||||
auto per_stage_rank_num = device_num / stage_num;
|
||||
return rank_id / per_stage_rank_num;
|
||||
}
|
||||
|
||||
// Only auto_parallel and semi_auto_parallel support PipelineSplit
|
||||
bool PipelineSplit(const ResourcePtr &res) {
|
||||
auto parallel_mode = parallel::ParallelContext::GetInstance()->parallel_mode();
|
||||
if (parallel_mode != parallel::SEMI_AUTO_PARALLEL || parallel_mode != parallel::AUTO_PARALLEL) {
|
||||
MS_LOG(INFO) << "Only auto_parallel and semi_auto_parallel support pipeline split.";
|
||||
return true;
|
||||
}
|
||||
auto stage_num = parallel::ParallelContext::GetInstance()->pipeline_stage_split_num();
|
||||
if (stage_num <= 1) {
|
||||
MS_LOG(INFO) << "stage num is: " << stage_num << ". No need Pipeline split.";
|
||||
return true;
|
||||
}
|
||||
auto manager = res->manager();
|
||||
auto root = res->func_graph();
|
||||
auto global_rank = GetRank();
|
||||
auto device_num = parallel::ParallelContext::GetInstance()->device_num();
|
||||
auto stage = InferStage(global_rank, stage_num, device_num);
|
||||
auto per_stage_rank_num = device_num / stage_num;
|
||||
auto transformer =
|
||||
std::make_shared<parallel::PipelineTransformer>(manager, stage, root, global_rank, per_stage_rank_num);
|
||||
// step1: Do color graph
|
||||
transformer->Coloring();
|
||||
// step2: Do color broadcast
|
||||
transformer->BroadCastColoring();
|
||||
// step3: Handle shared parameters
|
||||
transformer->ParameterColoring();
|
||||
transformer->HandleSharedParameter();
|
||||
// step4: Cut Graph
|
||||
transformer->CutGraph();
|
||||
// step5: Handle Sens
|
||||
transformer->CoverSensShape();
|
||||
// step6: Elim Graph stages and no used parameter
|
||||
transformer->ElimGraphStage();
|
||||
transformer->ElimParameter();
|
||||
return true;
|
||||
}
|
||||
} // namespace pipeline
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,28 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_CCSRC_PIPELINE_JIT_PIPELINE_SPLIT_H_
|
||||
#define MINDSPORE_CCSRC_PIPELINE_JIT_PIPELINE_SPLIT_H_
|
||||
|
||||
#include "pipeline/jit/resource.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace pipeline {
|
||||
bool PipelineSplit(const ResourcePtr &res);
|
||||
} // namespace pipeline
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_PIPELINE_JIT_PIPELINE_SPLIT_H_
|
||||
Loading…
Reference in new issue