mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-28 20:11:22 +00:00
to allow parallel training with mpi4py (#4942)
to allow parallel training with mpi4py Co-authored-by: liqun <liqun@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
liqunfu
GitHub
parent
9388d49c0d
commit
bb13b52291
3 changed files with 27 additions and 9 deletions
|
|
@ -235,7 +235,9 @@ void addObjectMethodsForTraining(py::module& m) {
|
|||
OrtPybindThrowIfError(sess->GetSessionHandle()->Load(path));
|
||||
|
||||
#if defined(USE_NCCL)
|
||||
CopyMPIContextToTrainingParameters(parameters, sess->GetSessionHandle()->GetLogger());
|
||||
bool use_nccl = parameters.allreduce_post_accumulation;
|
||||
if (!use_nccl && parameters.world_size > 1)
|
||||
CopyMPIContextToTrainingParameters(parameters, sess->GetSessionHandle()->GetLogger());
|
||||
#endif
|
||||
const auto config_result = ConfigureSessionForTraining(static_cast<TrainingSession*>(sess->GetSessionHandle()), parameters);
|
||||
|
||||
|
|
@ -249,7 +251,9 @@ void addObjectMethodsForTraining(py::module& m) {
|
|||
OrtPybindThrowIfError(sess->GetSessionHandle()->Load(buffer));
|
||||
|
||||
#if defined(USE_NCCL)
|
||||
CopyMPIContextToTrainingParameters(parameters, sess->GetSessionHandle()->GetLogger());
|
||||
bool use_nccl = parameters.allreduce_post_accumulation;
|
||||
if (!use_nccl && parameters.world_size > 1)
|
||||
CopyMPIContextToTrainingParameters(parameters, sess->GetSessionHandle()->GetLogger());
|
||||
#endif
|
||||
const auto config_result = ConfigureSessionForTraining(static_cast<TrainingSession*>(sess->GetSessionHandle()), parameters);
|
||||
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ class ORTGlueTest(unittest.TestCase):
|
|||
self.learning_rate = 2e-5
|
||||
self.num_train_epochs = 3.0
|
||||
self.local_rank = -1
|
||||
self.world_size = 1
|
||||
self.overwrite_output_dir = True
|
||||
self.gradient_accumulation_steps = 1
|
||||
self.data_dir = "/bert_data/hf_data/glue_data/"
|
||||
|
|
@ -279,7 +280,8 @@ class ORTGlueTest(unittest.TestCase):
|
|||
train_dataset=train_dataset,
|
||||
eval_dataset=eval_dataset,
|
||||
compute_metrics=compute_metrics,
|
||||
use_new_api=use_new_api
|
||||
use_new_api=use_new_api,
|
||||
world_size=self.world_size,
|
||||
)
|
||||
|
||||
# Training
|
||||
|
|
@ -303,17 +305,17 @@ class ORTGlueTest(unittest.TestCase):
|
|||
return results
|
||||
|
||||
if __name__ == "__main__":
|
||||
if get_mpi_context_world_size() > 1:
|
||||
local_rank = get_mpi_context_local_rank()
|
||||
world_size = get_mpi_context_world_size()
|
||||
if world_size > 1:
|
||||
# mpi launch
|
||||
logger.warning("mpirun launch, local_rank / world_size: %s : % s", local_rank, world_size)
|
||||
|
||||
print("mpirun launch")
|
||||
# TrainingArguments._setup_devices will call torch.distributed.init_process_group(backend="nccl")
|
||||
# pytorch expects following environment settings (which would be set if launched with torch.distributed.launch).
|
||||
|
||||
local_rank = get_mpi_context_local_rank()
|
||||
print("get_mpi_context_local_rank(): ", local_rank)
|
||||
os.environ['RANK'] = str(local_rank)
|
||||
os.environ['WORLD_SIZE'] = str(get_mpi_context_world_size())
|
||||
os.environ['WORLD_SIZE'] = str(world_size)
|
||||
os.environ['MASTER_ADDR'] = '127.0.0.1'
|
||||
os.environ['MASTER_PORT'] = '29500'
|
||||
|
||||
|
|
@ -323,6 +325,7 @@ if __name__ == "__main__":
|
|||
test = ORTGlueTest()
|
||||
test.setUp()
|
||||
test.local_rank = local_rank
|
||||
test.world_size = world_size
|
||||
test.test_bert_with_mrpc()
|
||||
else:
|
||||
unittest.main()
|
||||
|
|
|
|||
|
|
@ -106,6 +106,7 @@ class ORTTransformerTrainer:
|
|||
train_dataset: Dataset,
|
||||
eval_dataset: Dataset,
|
||||
compute_metrics: Callable[[EvalPrediction], Dict],
|
||||
world_size: Optional[int] = 1,
|
||||
use_new_api : Optional[bool] = False,
|
||||
):
|
||||
"""
|
||||
|
|
@ -115,6 +116,7 @@ class ORTTransformerTrainer:
|
|||
self.model_desc = model_desc
|
||||
self.new_model_desc = new_model_desc
|
||||
self.args = args
|
||||
self.world_size = world_size
|
||||
self.data_collator = DefaultDataCollator()
|
||||
self.train_dataset = train_dataset
|
||||
self.eval_dataset = eval_dataset
|
||||
|
|
@ -177,6 +179,7 @@ class ORTTransformerTrainer:
|
|||
|
||||
loss_scaler = amp.DynamicLossScaler() if self.args.fp16 else None
|
||||
device = self.args.device.type
|
||||
|
||||
device = f'{device}:{self.args.device.index}' if self.args.device.index else f'{device}:0'
|
||||
options = orttrainer.ORTTrainerOptions({'batch' : {
|
||||
'gradient_accumulation_steps' : self.args.gradient_accumulation_steps},
|
||||
|
|
@ -187,7 +190,13 @@ class ORTTransformerTrainer:
|
|||
'debug': {'deterministic_compute': True, },
|
||||
'utils': {
|
||||
'grad_norm_clip': False},
|
||||
'distributed': {'allreduce_post_accumulation': True},
|
||||
'distributed': {
|
||||
# we are running single node multi gpu test. thus world_rank = local_rank
|
||||
# and world_size = self.args.n_gpu
|
||||
'world_rank': max(0, self.args.local_rank),
|
||||
'world_size': int(self.world_size),
|
||||
'local_rank': max(0, self.args.local_rank),
|
||||
'allreduce_post_accumulation': True},
|
||||
'lr_scheduler': lr_scheduler
|
||||
})
|
||||
|
||||
|
|
@ -217,6 +226,8 @@ class ORTTransformerTrainer:
|
|||
learning_rate_description=IODescription('Learning_Rate', [1,], torch.float32),
|
||||
device=self.args.device,
|
||||
gradient_accumulation_steps=self.args.gradient_accumulation_steps,
|
||||
world_rank=max(0, self.args.local_rank),
|
||||
world_size=int(self.world_size),
|
||||
use_mixed_precision=self.args.fp16,
|
||||
allreduce_post_accumulation=True,
|
||||
get_lr_this_step=get_lr_this_step,
|
||||
|
|
|
|||
Loading…
Reference in a new issue