Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
@snnn
snnn committed Jun 12, 2024
1 parent ee284cb commit 75d75ef
Show file tree
Hide file tree
Showing 5 changed files with 23 additions and 22 deletions.
2 changes: 1 addition & 1 deletion orttraining/orttraining/core/framework/ortmodule_graph_builder.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ Status OrtModuleGraphBuilder::OptimizeForwardGraph(const TrainingGraphTransforme
}

if (!config.optimized_pre_grad_filepath.empty()) {
ORT_RETURN_IF_ERROR(Model::Save(*forward_model_, config.optimized_pre_grad_filepath));
ORT_RETURN_IF_ERROR(Model::Save(*forward_model_, ToPathString(config.optimized_pre_grad_filepath)));
}

return Status::OK();
Expand Down
1 change: 1 addition & 0 deletions orttraining/orttraining/core/optimizer/graph_transformer_config.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ struct TrainingGraphTransformerConfiguration : public GraphTransformerConfigurat
bool print_input_density{false};

// Path for serialization of the transformed optimized model. If empty, serialization is disabled.
// A UTF-8 string.
std::string optimized_pre_grad_filepath;
};

Expand Down
2 changes: 1 addition & 1 deletion orttraining/orttraining/python/orttraining_pybind_state.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -617,7 +617,7 @@ void addObjectMethodsForTraining(py::module& m) {
ORT_THROW_IF_ERROR(gradient_graph_builder->builder_->Build());
})
.def("save", [](PyGradientGraphBuilderContext* gradient_graph_builder, const std::string& path) {
ORT_THROW_IF_ERROR(Model::Save(*(gradient_graph_builder->model_), path));
ORT_THROW_IF_ERROR(Model::Save(*(gradient_graph_builder->model_), ToPathString(path)));
})
.def("get_model", [](PyGradientGraphBuilderContext* gradient_graph_builder) {
std::string model_str;
Expand Down
12 changes: 6 additions & 6 deletions orttraining/orttraining/test/gradient/allreduce_op_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -472,7 +472,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumAllreduceOptimizerTest) {
build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::GpuHierarchicalReduction, true /*build_optimizer*/,
false /*half_precision*/);

std::string model_file_name = "GPUHierarchicalAdasumAllreduceOptimizerTest.onnx";
PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumAllreduceOptimizerTest.onnx");
auto status = onnxruntime::Model::Save(model, model_file_name);

SessionOptions so;
Expand Down Expand Up @@ -649,7 +649,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumAllreduceOptimizerFP16Test) {
build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::GpuHierarchicalReduction, true /*build_optimizer*/,
true /*half_precision*/);

std::string model_file_name = "GPUHierarchicalAdasumAllreduceOptimizerFP16Test.onnx";
PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumAllreduceOptimizerFP16Test.onnx");
auto status = onnxruntime::Model::Save(model, model_file_name);

SessionOptions so;
Expand Down Expand Up @@ -791,7 +791,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumAllreduceTest) {
adasum_graph_configs.push_back(adasum_graph_config);
build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::GpuHierarchicalReduction);

std::string model_file_name = "GPUHierarchicalAdasumAllreduceTest.onnx";
PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumAllreduceTest.onnx");
auto status = onnxruntime::Model::Save(model, model_file_name);

SessionOptions so;
Expand Down Expand Up @@ -896,7 +896,7 @@ TEST(AllreduceTest, GPUHierarchicalAdasumFP16AllreduceTest) {
false /*build_optimizer*/,
true /*half_precision*/);

std::string model_file_name = "GPUHierarchicalAdasumFP16AllreduceTest.onnx";
PathString model_file_name = ORT_TSTR("GPUHierarchicalAdasumFP16AllreduceTest.onnx");
auto status = onnxruntime::Model::Save(model, model_file_name);

SessionOptions so;
Expand Down Expand Up @@ -1003,7 +1003,7 @@ TEST(AllreduceTest, GPUAdasumAllreduceTest) {

build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::CpuReduction);

std::string model_file_name = "GPUAdasumAllreduceTest.onnx";
PathString model_file_name = ORT_TSTR("GPUAdasumAllreduceTest.onnx");
auto status = onnxruntime::Model::Save(model, model_file_name);

SessionOptions so;
Expand Down Expand Up @@ -1110,7 +1110,7 @@ TEST(AllreduceTest, GPUAdasumFP16AllreduceTest) {

build_allreduce_graph(graph, adasum_graph_configs, training::AdasumReductionType::CpuReduction, true /*half_precision*/);

std::string model_file_name = "GPUAdasumFP16AllreduceTest.onnx";
PathString model_file_name = ORT_TSTR("GPUAdasumFP16AllreduceTest.onnx");
auto status = onnxruntime::Model::Save(model, model_file_name);

SessionOptions so;
Expand Down
28 changes: 14 additions & 14 deletions orttraining/orttraining/test/optimizer/graph_transform_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -627,7 +627,7 @@ TEST_F(GraphTransformationTests, MegatronMLPPartitionRank0) {
TransformerLevel::Level1));
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));

auto model_uri2 = "mlp_megatron_basic_test_partition_rank0.onnx";
PathString model_uri2 = ORT_TSTR("mlp_megatron_basic_test_partition_rank0.onnx");
ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2));

{
Expand Down Expand Up @@ -705,7 +705,7 @@ TEST_F(GraphTransformationTests, MegatronMLPPartitionRank1) {
TransformerLevel::Level1));
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));

auto model_uri2 = "mlp_megatron_basic_test_partition_rank1.onnx";
PathString model_uri2 = ORT_TSTR("mlp_megatron_basic_test_partition_rank1.onnx");
ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2));

{
Expand Down Expand Up @@ -765,7 +765,7 @@ TEST_F(GraphTransformationTests, MegatronMLPPartitionRank1) {
}

TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank0) {
auto model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx";
PathString model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx";
std::shared_ptr<Model> p_model;
ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
Graph& graph = p_model->MainGraph();
Expand All @@ -781,7 +781,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank0) {
TransformerLevel::Level1));
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));

auto model_uri2 = "self_attention_megatron_basic_test_partition_rank0.onnx";
PathString model_uri2 = ORT_TSTR("self_attention_megatron_basic_test_partition_rank0.onnx");
ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2));

{
Expand Down Expand Up @@ -838,7 +838,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank0) {
}

TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank1) {
auto model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx";
PathString model_uri = MODEL_FOLDER "model_parallel/self_attention_megatron_basic_test.onnx";
std::shared_ptr<Model> p_model;
ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
Graph& graph = p_model->MainGraph();
Expand All @@ -856,7 +856,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank1) {
TransformerLevel::Level1));
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));

auto model_uri2 = "self_attention_megatron_basic_test_partition_rank1.onnx";
PathString model_uri2 = ORT_TSTR("self_attention_megatron_basic_test_partition_rank1.onnx");
ASSERT_STATUS_OK(Model::Save(*p_model, model_uri2));

{
Expand Down Expand Up @@ -913,7 +913,7 @@ TEST_F(GraphTransformationTests, MegatronSelfAttentionPartitionRank1) {
}

TEST_F(GraphTransformationTests, BiasGeluRecomputeTest) {
auto model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion_recompute.onnx";
PathString model_uri = MODEL_FOLDER "fusion/bias_gelu_fusion_recompute.onnx";
std::shared_ptr<Model> p_model;
ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_));
Graph& graph = p_model->MainGraph();
Expand Down Expand Up @@ -1397,15 +1397,15 @@ static void RunPartitionCorrectnessTest(std::string model_path,
TransformerLevel::Level1));
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, logger));
graphs.push_back(&graph);
auto model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_rank_") + ToPathString(std::to_string(i)) + ORT_TSTR(".onnx");
PathString model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_rank_") + ToPathString(std::to_string(i)) + ORT_TSTR(".onnx");
ASSERT_STATUS_OK(Model::Save(*p_models[i], model_uri2));
}

onnxruntime::Model combine_model("combine_graph", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 12}, {kMSDomain, 1}}, {}, logger);
auto& combine_graph = combine_model.MainGraph();
auto ret = horizontal_parallel_test_utils::MergeGraphsOnAllWorkers(graphs, combine_graph);
ORT_ENFORCE(ret.IsOK());
auto model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_combine.onnx");
PathString model_uri2 = ToPathString(model_path) + ORT_TSTR("_partition_combine.onnx");
ASSERT_STATUS_OK(Model::Save(combine_model, model_uri2));

float scale = 1.f;
Expand Down Expand Up @@ -1790,7 +1790,7 @@ TEST_F(GraphTransformationTests, ScaledSumFusionTwoInputs) {

#ifdef ENABLE_TRITON
TEST_F(GraphTransformationTests, TritonFusion) {
auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
Graph& graph = model->MainGraph();
Expand All @@ -1805,7 +1805,7 @@ TEST_F(GraphTransformationTests, TritonFusion) {
ASSERT_TRUE(op_to_count["LayerNormalization"] == 4);

{
auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
Graph& graph = model->MainGraph();
Expand Down Expand Up @@ -1845,7 +1845,7 @@ TEST_F(GraphTransformationTests, TritonFusion) {

// No Dropout.
{
auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
Graph& graph = model->MainGraph();
Expand Down Expand Up @@ -1884,7 +1884,7 @@ TEST_F(GraphTransformationTests, TritonFusion) {

// Ignore min nodes.
{
auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
Graph& graph = model->MainGraph();
Expand Down Expand Up @@ -1924,7 +1924,7 @@ TEST_F(GraphTransformationTests, TritonFusion) {

// Exclude Softmax using axis attribute.
{
auto model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
PathString model_uri = MODEL_FOLDER "bert_toy_opset14.onnx";
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
Graph& graph = model->MainGraph();
Expand Down

0 comments on commit 75d75ef

Please sign in to comment.