diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index 7e5b2e81a3..3cc0e46546 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -109,6 +109,19 @@ bool IsRootNode(const TrainingSession::TrainingConfiguration& config) { } } // namespace + +void TrainingSession::FilterUnusedWeights(const std::unordered_set& weight_names_to_train, + std::unordered_set& filtered_weight_names_to_train) { + filtered_weight_names_to_train.clear(); + for (const auto& name: weight_names_to_train) { + auto nodes = model_->MainGraph().GetConsumerNodes(name); + if (!nodes.empty()) + filtered_weight_names_to_train.insert(name); + else + LOGS(*session_logger_, WARNING) << "Couldn't find any consumer node for weight " << name << ", exclude it from training."; + } +} + Status TrainingSession::ConfigureForTraining( const TrainingConfiguration& config, TrainingConfigurationResult& config_result_out) { ORT_RETURN_IF( @@ -117,6 +130,9 @@ Status TrainingSession::ConfigureForTraining( if (is_configured_) return Status::OK(); + std::unordered_set filtered_config_weight_names_to_train; + FilterUnusedWeights(config.weight_names_to_train, filtered_config_weight_names_to_train); + TrainingConfigurationResult config_result{}; ORT_ENFORCE(config.distributed_config.pipeline_parallel_size > 0, @@ -186,8 +202,8 @@ Status TrainingSession::ConfigureForTraining( // For case we use GetTrainableModelInitializers to get trainable weights such as C++ frontend, it may get more initializers // than trainable weights here as it's before transformers. So the constant folding may miss some nodes we actually can fold. std::unordered_set trainable_initializers = - !config.weight_names_to_train.empty() - ? config.weight_names_to_train + !filtered_config_weight_names_to_train.empty() + ? filtered_config_weight_names_to_train : GetTrainableModelInitializers(config.immutable_weights, loss_name); if (config.weight_names_to_not_train.size() > 0) { LOGS(*session_logger_, INFO) << "Excluding following weights from trainable list as specified in configuration:"; @@ -201,8 +217,8 @@ Status TrainingSession::ConfigureForTraining( // derive actual set of weights to train std::unordered_set weight_names_to_train = - !config.weight_names_to_train.empty() - ? config.weight_names_to_train + !filtered_config_weight_names_to_train.empty() + ? filtered_config_weight_names_to_train : GetTrainableModelInitializers(config.immutable_weights, loss_name); for (const auto& weight_name_to_not_train : config.weight_names_to_not_train) { weight_names_to_train.erase(weight_name_to_not_train); diff --git a/orttraining/orttraining/core/session/training_session.h b/orttraining/orttraining/core/session/training_session.h index feb17d7f7a..cbd1d20aef 100644 --- a/orttraining/orttraining/core/session/training_session.h +++ b/orttraining/orttraining/core/session/training_session.h @@ -450,6 +450,9 @@ class TrainingSession : public InferenceSession { NameMLValMap GetWeights() const; + void FilterUnusedWeights(const std::unordered_set& weight_names_to_train, + std::unordered_set& filtered_weight_names_to_train); + static bool IsImmutableWeight(const ImmutableWeights& immutable_weights, const Node* node, const TensorProto* weight_tensor, diff --git a/orttraining/orttraining/test/python/orttraining_run_glue.py b/orttraining/orttraining/test/python/orttraining_run_glue.py index fb9a6476ad..a5d196cc69 100644 --- a/orttraining/orttraining/test/python/orttraining_run_glue.py +++ b/orttraining/orttraining/test/python/orttraining_run_glue.py @@ -67,11 +67,35 @@ class ORTGlueTest(unittest.TestCase): self.cache_dir = '/tmp/glue/' self.logging_steps = 10 + def test_roberta_with_mrpc(self): + expected_acc = 0.8897058823529411 + expected_f1 = 0.9200710479573712 + expected_acc_and_f1 = 0.9048884651551561 + expected_loss = 0.2911236987394445 + + results = self.run_glue(model_name="roberta-base", task_name="MRPC", fp16=False) + assert_allclose(results['acc'], expected_acc) + assert_allclose(results['f1'], expected_f1) + assert_allclose(results['acc_and_f1'], expected_acc_and_f1) + assert_allclose(results['loss'], expected_loss) + + def test_roberta_fp16_with_mrpc(self): + expected_acc = 0.8921568627450981 + expected_f1 = 0.9219858156028369 + expected_acc_and_f1 = 0.9070713391739675 + expected_loss = 0.3033953265232198 + + results = self.run_glue(model_name="roberta-base", task_name="MRPC", fp16=True) + assert_allclose(results['acc'], expected_acc) + assert_allclose(results['f1'], expected_f1) + assert_allclose(results['acc_and_f1'], expected_acc_and_f1) + assert_allclose(results['loss'], expected_loss) + def test_bert_with_mrpc(self): - expected_acc = 0.8578431372549019 - expected_f1 = 0.9003436426116839 - expected_acc_and_f1 = 0.8790933899332929 - expected_loss = 0.415903969430456 + expected_acc = 0.8529411764705882 + expected_f1 = 0.896551724137931 + expected_acc_and_f1 = 0.8747464503042597 + expected_loss = 0.4139287974320206 results = self.run_glue(model_name="bert-base-cased", task_name="MRPC", fp16=False) assert_allclose(results['acc'], expected_acc) @@ -80,10 +104,10 @@ class ORTGlueTest(unittest.TestCase): assert_allclose(results['loss'], expected_loss) def test_bert_fp16_with_mrpc(self): - expected_acc = 0.8529411764705882 - expected_f1 = 0.8951048951048952 - expected_acc_and_f1 = 0.8740230357877417 - expected_loss = 0.36075809042827756 + expected_acc = 0.8627450980392157 + expected_f1 = 0.9047619047619047 + expected_acc_and_f1 = 0.8837535014005602 + expected_loss = 0.41143255315574945 results = self.run_glue(model_name="bert-base-cased", task_name="MRPC", fp16=True) assert_allclose(results['acc'], expected_acc) @@ -91,11 +115,32 @@ class ORTGlueTest(unittest.TestCase): assert_allclose(results['acc_and_f1'], expected_acc_and_f1) assert_allclose(results['loss'], expected_loss) + def model_to_desc(self, model_name, model): + if model_name.startswith('bert') or model_name.startswith('xlnet'): + model_desc = ModelDescription([ + IODescription('input_ids', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=model.config.vocab_size), + IODescription('attention_mask', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=2), + IODescription('token_type_ids', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=2), + IODescription('labels', ['batch',], torch.int64, num_classes=2)], [ + IODescription('loss', [], torch.float32), + IODescription('logits', ['batch', 2], torch.float32)]) + elif model_name.startswith('roberta'): + model_desc = ModelDescription([ + IODescription('input_ids', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=model.config.vocab_size), + IODescription('attention_mask', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=2), + IODescription('labels', ['batch',], torch.int64, num_classes=2)], [ + IODescription('loss', [], torch.float32), + IODescription('logits', ['batch', 2], torch.float32)]) + else: + raise RuntimeError("unsupported base model name {}.".format(model_name)) + + return model_desc + def run_glue(self, model_name, task_name, fp16): model_args = ModelArguments(model_name_or_path=model_name, cache_dir=self.cache_dir) data_args = GlueDataTrainingArguments(task_name=task_name, data_dir=self.data_dir + "/" + task_name, max_seq_length=self.max_seq_length) - + training_args = TrainingArguments(output_dir=self.output_dir + "/" + task_name, do_train=True, do_eval=True, per_gpu_train_batch_size=self.train_batch_size, learning_rate=self.learning_rate, num_train_epochs=self.num_train_epochs,local_rank=self.local_rank, @@ -164,14 +209,7 @@ class ORTGlueTest(unittest.TestCase): preds = np.squeeze(p.predictions) return glue_compute_metrics(data_args.task_name, preds, p.label_ids) - model_desc = ModelDescription([ - IODescription('input_ids', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=model.config.vocab_size), - IODescription('attention_mask', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=2), - IODescription('token_type_ids', ['batch', 'max_seq_len_in_batch'], torch.int64, num_classes=2), - IODescription('labels', ['batch',], torch.int64, num_classes=2)], [ - IODescription('loss', [], torch.float32), - IODescription('logits', ['batch', 2], torch.float32)]) - + model_desc = self.model_to_desc(model_name, model) # Initialize the ORTTrainer within ORTTransformerTrainer trainer = ORTTransformerTrainer( model=model,