def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
megatron_amp_o2 = cfg.model.get('megatron_amp_O2', False)
plugins = [
NLPDDPPlugin(
no_ddp_communication_hook=(megatron_amp_o2
and cfg.trainer.precision == 'bf16'
), # Only bf16 uses fp32_grad_accum.
gradient_as_bucket_view=cfg.model.gradient_as_bucket_view,
find_unused_parameters=False,
)
]
if cfg.trainer.precision in [16, 'bf16']:
scaler = None
if cfg.trainer.precision == 16:
scaler = GradScaler(
init_scale=cfg.model.get('native_amp_init_scale', 2**32),
growth_interval=cfg.model.get('native_amp_growth_interval',
1000),
hysteresis=cfg.model.get('hysteresis', 2),
)
if megatron_amp_o2:
plugins.append(
MegatronHalfPrecisionPlugin(precision=cfg.trainer.precision,
device='cuda',
scaler=scaler))
else:
plugins.append(
NativeMixedPrecisionPlugin(precision=cfg.trainer.precision,
device='cuda',
scaler=scaler))
if cfg.get('cluster_type', None) == 'BCP':
plugins.append(TorchElasticEnvironment())
trainer = Trainer(plugins=plugins, **cfg.trainer)
exp_manager(trainer, cfg.exp_manager)
# update resume from checkpoint found by exp_manager
resume_from_checkpoint = trainer.checkpoint_connector.resume_from_checkpoint_fit_path
logging.info(
f'Resuming training from checkpoint: {resume_from_checkpoint}')
trainer.checkpoint_connector = CheckpointConnector(
trainer, resume_from_checkpoint=resume_from_checkpoint)
# Override timer callback to a stateless one
for idx, callback in enumerate(trainer.callbacks):
if isinstance(callback, Timer):
trainer.callbacks[idx] = StatelessTimer(cfg.trainer.max_time, )
# hydra interpolation does not work here as the interpolation key is lost when PTL saves hparams
with open_dict(cfg):
cfg.model.precision = cfg.trainer.precision
model = MegatronXNlIModel(cfg.model, trainer)
trainer.fit(model)
trainer.test(model)
def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
plugins = [NLPDDPPlugin(num_nodes=cfg.trainer.num_nodes)]
if cfg.trainer.precision == 16:
scaler = GradScaler(
init_scale=cfg.model.get('native_amp_init_scale', 2**32),
growth_interval=cfg.model.get('native_amp_growth_interval', 1000),
)
plugins.append(
NativeMixedPrecisionPlugin(precision=16,
device='cuda',
scaler=scaler))
if cfg.get('cluster_type', None) == 'BCP':
plugins.append(TorchElasticEnvironment())
trainer = Trainer(plugins=plugins, **cfg.trainer)
exp_manager(trainer, cfg.exp_manager)
# update resume from checkpoint found by exp_manager
resume_from_checkpoint = trainer.checkpoint_connector.resume_from_checkpoint_fit_path
if resume_from_checkpoint is not None:
# inject mp_rank into resume_from_checkpoint
if cfg.model.tensor_model_parallel_size is not None and cfg.model.tensor_model_parallel_size > 1:
mp_rank = compute_model_parallel_rank(
trainer.local_rank, cfg.model.tensor_model_parallel_size)
resume_from_checkpoint = Path(resume_from_checkpoint)
resume_from_checkpoint = resume_from_checkpoint.parent.parent.joinpath(
f'mp_rank_{mp_rank:02d}').joinpath(resume_from_checkpoint.name)
resume_from_checkpoint = str(resume_from_checkpoint)
logging.info(
f'Resuming training from checkpoint: {resume_from_checkpoint}')
trainer.checkpoint_connector = CheckpointConnector(
trainer, resume_from_checkpoint=resume_from_checkpoint)
# Override timer callback to a stateless one
for idx, callback in enumerate(trainer.callbacks):
if isinstance(callback, Timer):
trainer.callbacks[idx] = StatelessTimer(cfg.trainer.max_time, )
# hydra interpolation does not work here as the interpolation key is lost when PTL saves hparams
with open_dict(cfg):
cfg.model.precision = cfg.trainer.precision
model = MegatronT5Model(cfg.model, trainer)
trainer.fit(model)
def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
plugins = [NLPDDPPlugin(num_nodes=cfg.trainer.num_nodes)]
if cfg.trainer.precision == 16:
plugins.append(
NLPNativeMixedPrecisionPlugin(
init_scale=cfg.model.get('native_amp_init_scale', 2**32),
growth_interval=cfg.model.get('native_amp_growth_interval',
1000),
))
elif cfg.trainer.precision == 'bf16':
plugins.append(NLPNativeBfloat16PrecisionPlugin())
else:
plugins.append(NLPPrecisionPlugin())
if cfg.get('cluster_type', None) == 'BCP':
plugins.append(TorchElasticEnvironment())
trainer = Trainer(plugins=plugins, **cfg.trainer)
exp_manager(trainer, cfg.exp_manager)
# update resume from checkpoint found by exp_manager
resume_from_checkpoint = trainer.resume_from_checkpoint
if resume_from_checkpoint is not None:
mp_rank = compute_model_parallel_rank(
trainer.local_rank, cfg.model.tensor_model_parallel_size)
resume_from_checkpoint = Path(resume_from_checkpoint)
resume_from_checkpoint = resume_from_checkpoint.parent.parent.joinpath(
f'mp_rank_{mp_rank:02d}').joinpath(resume_from_checkpoint.name)
resume_from_checkpoint = str(resume_from_checkpoint)
logging.info(
f'Resuming training from checkpoint: {resume_from_checkpoint}')
trainer.checkpoint_connector = CheckpointConnector(
trainer, resume_from_checkpoint=resume_from_checkpoint)
# Override timer callback to a stateless one
for idx, callback in enumerate(trainer.callbacks):
if isinstance(callback, Timer):
trainer.callbacks[idx] = StatelessTimer(cfg.trainer.max_time, )
model = MegatronGPTModel(cfg.model, trainer)
trainer.fit(model)
def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
plugins = [NLPDDPPlugin(find_unused_parameters=False)]
if cfg.trainer.precision == 16:
scaler = GradScaler(
init_scale=cfg.model.get('native_amp_init_scale', 2**32),
growth_interval=cfg.model.get('native_amp_growth_interval', 1000),
)
plugins.append(
NativeMixedPrecisionPlugin(precision=16,
device='cuda',
scaler=scaler))
if cfg.get('cluster_type', None) == 'BCP':
plugins.append(TorchElasticEnvironment())
trainer = Trainer(plugins=plugins, **cfg.trainer)
exp_manager(trainer, cfg.exp_manager)
# update resume from checkpoint found by exp_manager
resume_from_checkpoint = trainer.checkpoint_connector.resume_from_checkpoint_fit_path
logging.info(
f'Resuming training from checkpoint: {resume_from_checkpoint}')
trainer.checkpoint_connector = CheckpointConnector(
trainer, resume_from_checkpoint=resume_from_checkpoint)
# Override timer callback to a stateless one
for idx, callback in enumerate(trainer.callbacks):
if isinstance(callback, Timer):
trainer.callbacks[idx] = StatelessTimer(cfg.trainer.max_time, )
# hydra interpolation does not work here as the interpolation key is lost when PTL saves hparams
with open_dict(cfg):
cfg.model.precision = cfg.trainer.precision
model = MegatronBertModel(cfg.model, trainer)
trainer.fit(model)
def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
# setup the data processor
for processor_config in cfg.model.task_processors:
processor = TemplateProcessor(
template=processor_config.template, limit_length_field=processor_config.limit_length_field
)
register_taskdata_processor(processor_config.taskname, processor)
plugins = [NLPDDPPlugin()]
if cfg.trainer.precision == 16:
scaler = GradScaler(
init_scale=cfg.model.get('native_amp_init_scale', 2 ** 32),
growth_interval=cfg.model.get('native_amp_growth_interval', 1000),
)
plugins.append(NativeMixedPrecisionPlugin(precision=16, device='cuda', scaler=scaler))
if cfg.get('cluster_type', None) == 'BCP':
plugins.append(TorchElasticEnvironment())
trainer = Trainer(plugins=plugins, **cfg.trainer)
exp_manager(trainer, cfg.exp_manager)
# update resume from checkpoint found by exp_manager
resume_from_checkpoint = trainer.checkpoint_connector.resume_from_checkpoint_fit_path
logging.info(f'Resuming training from checkpoint: {resume_from_checkpoint}')
trainer.checkpoint_connector = CheckpointConnector(trainer, resume_from_checkpoint=resume_from_checkpoint)
# Override timer callback to a stateless one
for idx, callback in enumerate(trainer.callbacks):
if isinstance(callback, Timer):
trainer.callbacks[idx] = StatelessTimer(cfg.trainer.max_time,)
# hydra interpolation does not work here as the interpolation key is lost when PTL saves hparams
with open_dict(cfg):
cfg.model.precision = cfg.trainer.precision
model = MegatronGPTPTuneModel(cfg.model, trainer)
trainer.fit(model)
if cfg.model.data.test_ds.file_path:
logging.info("===========================================================================================")
logging.info("Starting the testing of the trained model on test set...")
trainer.test(model)
logging.info("Testing finished!")
logging.info("===========================================================================================")
# extract the path of the best checkpoint from the training, you may update it to any checkpoint
checkpoint_path = trainer.checkpoint_callback.best_model_path
tensor_parallel_size = cfg.model.tensor_model_parallel_size
pathobj = pathlib.Path(checkpoint_path)
checkpoint_folder = str(pathobj.parent)
checkpoint_name = str(pathobj.name)
rank = trainer.accelerator.training_type_plugin.local_rank
if tensor_parallel_size > 1:
# inject model parallel rank
checkpoint_path = os.path.join(checkpoint_folder, f'mp_rank_{rank:02d}', checkpoint_name)
else:
checkpoint_path = os.path.join(checkpoint_folder, checkpoint_name)
# Load the checkpoint
best_eval_model = MegatronGPTPTuneModel.load_from_checkpoint(
checkpoint_path=checkpoint_path, strict=False, trainer=trainer
)
logging.info(f'Best checkpoint path: {checkpoint_path}')
logging.info("Running Test with best EVAL checkpoint!")
# setup the test dataset
# best_eval_model.setup_test_data(test_data_config=cfg.model.data.test_ds)
if torch.distributed.is_initialized():
torch.distributed.barrier()
trainer.test(model=best_eval_model, ckpt_path=None, verbose=False)
logging.info("Beset EVAL Testing finished!")
logging.info("===========================================================================================")
if cfg.model.nemo_path:
# '.nemo' file contains the last checkpoint and the params to initialize the model
best_eval_model.save_to(cfg.model.nemo_path)
logging.info(f'Model is saved into `.nemo` file: {cfg.model.nemo_path}')
# perform inference on a list of queries.
if "infer_samples" in cfg.model and cfg.model.infer_samples:
logging.info("===========================================================================================")
logging.info("Starting the inference on some sample queries...")
# max_seq_length=512 is the maximum length BERT supports.
results = best_eval_model.cuda().ptune_inference(
queries=cfg.model.infer_samples, batch_size=1, decode_token_len=5
)
logging.info('The prediction results of some sample queries with the trained model:')
for query, result in zip(cfg.model.infer_samples, results):
logging.info(f'Query : {query}')
logging.info(f'Predicted label: {result}')
logging.info("Inference finished!")
logging.info("===========================================================================================")
