def __init__(self, cfg: DictConfig, trainer: Trainer = None):
super().__init__(cfg, trainer)
# handles model parallel save and restore logic
self._save_restore_connector = NLPSaveRestoreConnector()
self.set_world_size(trainer)
if not HAVE_APEX:
logging.warning("Apex was not found. Using model parallel or megatron models will error out.")
def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
trainer = None
if cfg.trainer.precision == 16:
trainer = Trainer(
plugins=[
NLPDDPPlugin(),
NLPNativeMixedPrecisionPlugin(
init_scale=cfg.model.get('native_amp_init_scale', 2 ** 32),
growth_interval=cfg.model.get('native_amp_growth_interval', 1000),
),
],
**cfg.trainer,
)
elif cfg.trainer.precision == 'bf16':
trainer = Trainer(plugins=[NLPDDPPlugin(), NLPNativeBfloat16PrecisionPlugin(),], **cfg.trainer,)
else:
trainer = Trainer(plugins=[NLPDDPPlugin(), NLPPrecisionPlugin()], **cfg.trainer)
app_state = AppState()
app_state.model_parallel_size = cfg.model.tensor_model_parallel_size
app_state.model_parallel_rank = compute_model_parallel_rank(trainer.local_rank, app_state.model_parallel_size)
model = MegatronGPTModel.restore_from(
cfg.restore_from_path, trainer=trainer, save_restore_connector=NLPSaveRestoreConnector(),
)
# Note: most nemo models must have the data paths configured before instantiating the model
# MegatronGPTMOdel sets up the data in the PTL method .setup which happens after DDP spawns.
model.cfg.data.splits_string = cfg.model.data.splits_string
trainer.test(model)
def convert(rank, world_size, args):
app_state = AppState()
app_state.data_parallel_rank = 0
trainer = Trainer(gpus=args.tensor_model_parallel_size)
# TODO: reach out to PTL For an API-safe local rank override
trainer.accelerator.training_type_plugin._local_rank = rank
if args.tensor_model_parallel_size is not None and args.tensor_model_parallel_size > 1:
# inject model parallel rank
checkpoint_path = os.path.join(args.checkpoint_folder,
f'mp_rank_{rank:02d}',
args.checkpoint_name)
else:
checkpoint_path = os.path.join(args.checkpoint_folder,
args.checkpoint_name)
if args.model_type == 'gpt':
model = MegatronGPTModel.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
elif args.model_type == 'bert':
model = MegatronBertModel.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
elif args.model_type == 't5':
model = MegatronT5Model.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
model._save_restore_connector = NLPSaveRestoreConnector()
if torch.distributed.is_initialized():
torch.distributed.barrier()
model.save_to(args.nemo_file_path)
logging.info(f'NeMo model saved to: {args.nemo_file_path}')
def convert(local_rank, rank, world_size, args):
app_state = AppState()
app_state.data_parallel_rank = 0
num_nodes = world_size // args.gpus_per_node
if args.bcp:
trainer = Trainer(devices=args.gpus_per_node,
num_nodes=num_nodes,
accelerator='gpu',
plugins=[TorchElasticEnvironment()])
else:
trainer = Trainer(devices=args.gpus_per_node,
num_nodes=num_nodes,
accelerator='gpu')
app_state.pipeline_model_parallel_size = args.pipeline_model_parallel_size
app_state.tensor_model_parallel_size = args.tensor_model_parallel_size
app_state.model_parallel_size = app_state.tensor_model_parallel_size * app_state.pipeline_model_parallel_size
parallel_state.initialize_model_parallel(
tensor_model_parallel_size_=app_state.tensor_model_parallel_size,
pipeline_model_parallel_size_=app_state.pipeline_model_parallel_size,
)
app_state.pipeline_model_parallel_rank = parallel_state.get_pipeline_model_parallel_rank(
)
app_state.tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank(
)
# inject model parallel rank
checkpoint_path = inject_model_parallel_rank(
os.path.join(args.checkpoint_folder, args.checkpoint_name))
logging.info(
f'rank: {rank}, local_rank: {local_rank}, is loading checkpoint: {checkpoint_path} for tp_rank: {app_state.tensor_model_parallel_rank} and pp_rank: {app_state.pipeline_model_parallel_rank}'
)
if args.model_type == 'gpt':
model = MegatronGPTModel.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
elif args.model_type == 'bert':
model = MegatronBertModel.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
elif args.model_type == 't5':
model = MegatronT5Model.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
elif args.model_type == 'nmt':
model = MegatronNMTModel.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
model._save_restore_connector = NLPSaveRestoreConnector()
if torch.distributed.is_initialized():
torch.distributed.barrier()
model.save_to(args.nemo_file_path)
logging.info(f'NeMo model saved to: {args.nemo_file_path}')
def main(cfg) -> None:
# trainer required for restoring model parallel models
trainer = Trainer(plugins=NLPDDPPlugin(), **cfg.trainer)
assert (
cfg.trainer.devices * cfg.trainer.num_nodes
== cfg.tensor_model_parallel_size * cfg.pipeline_model_parallel_size
), "devices * num_nodes should equal tensor_model_parallel_size * pipeline_model_parallel_size"
app_state = AppState()
app_state.model_parallel_size = cfg.tensor_model_parallel_size * cfg.pipeline_model_parallel_size
(
app_state.tensor_model_parallel_rank,
app_state.pipeline_model_parallel_rank,
app_state.model_parallel_size,
app_state.data_parallel_size,
app_state.pipeline_model_parallel_split_rank,
) = fake_initialize_model_parallel(
world_size=app_state.model_parallel_size,
rank=trainer.global_rank,
tensor_model_parallel_size_=cfg.tensor_model_parallel_size,
pipeline_model_parallel_size_=cfg.pipeline_model_parallel_size,
pipeline_model_parallel_split_rank_=cfg.pipeline_model_parallel_split_rank,
)
if cfg.model_file is not None:
if not os.path.exists(cfg.model_file):
raise ValueError(f"Model file {cfg.model_file} does not exist")
model = MegatronNMTModel.restore_from(
restore_path=cfg.model_file, trainer=trainer, save_restore_connector=NLPSaveRestoreConnector(),
)
elif cfg.checkpoint_dir is not None:
checkpoint_path = inject_model_parallel_rank(os.path.join(cfg.checkpoint_dir, cfg.checkpoint_name))
model = MegatronNMTModel.load_from_checkpoint(checkpoint_path, hparams_file=cfg.hparams_file, trainer=trainer)
else:
raise ValueError("need at least a nemo file or checkpoint dir")
model.freeze()
logging.info(f"Translating: {cfg.srctext}")
src_text = []
translations = []
with open(cfg.srctext, 'r') as src_f, open(cfg.tgtout, 'w') as tgt_f:
for line in src_f:
src_text.append(line.strip())
if len(src_text) == cfg.batch_size:
translations = model.translate(
text=src_text, source_lang=cfg.source_lang, target_lang=cfg.target_lang,
)
for translation in translations:
tgt_f.write(translation + "\n")
src_text = []
if len(src_text) > 0:
translations = model.translate(text=src_text, source_lang=cfg.source_lang, target_lang=cfg.target_lang,)
for translation in translations:
tgt_f.write(translation + "\n")
def convert(rank, world_size, args):
app_state = AppState()
app_state.data_parallel_rank = 0
trainer = Trainer(gpus=args.tensor_model_parallel_size)
# TODO: reach out to PTL For an API-safe local rank override
trainer.accelerator.training_type_plugin._local_rank = rank
if args.tensor_model_parallel_size is not None and args.tensor_model_parallel_size > 1:
# inject model parallel rank
checkpoint_path = os.path.join(args.checkpoint_folder,
f'mp_rank_{rank:02d}',
args.checkpoint_name)
else:
checkpoint_path = os.path.join(args.checkpoint_folder,
args.checkpoint_name)
if args.model_type == 'gpt':
## this dictionary is used to rename the model parameters
name_translate = {}
name_translate['transformer'] = 'encoder'
name_translate['.attention.'] = '.self_attention.'
model = load_from_checkpoint(
MegatronGPTModel,
checkpoint_path,
hparams_file=args.hparams_file,
trainer=trainer,
translator=name_translate,
strict=False,
)
elif args.model_type == 'bert':
## this dictionary is used to rename the model parameters
name_translate = {}
name_translate['transformer'] = 'encoder'
name_translate['.attention.'] = '.self_attention.'
model = load_from_checkpoint(
MegatronBertModel,
checkpoint_path,
hparams_file=args.hparams_file,
trainer=trainer,
translator=name_translate,
strict=False,
)
else:
raise NotImplemented("{} is not supported".format(args.model_type))
model._save_restore_connector = NLPSaveRestoreConnector()
if torch.distributed.is_initialized():
torch.distributed.barrier()
model.save_to(args.nemo_file_path)
logging.info(f'NeMo model saved to: {args.nemo_file_path}')
def convert(rank, world_size, args):
app_state = AppState()
app_state.data_parallel_rank = 0
trainer = Trainer(gpus=args.tensor_model_parallel_size)
# TODO: reach out to PTL For an API-safe local rank override
trainer.accelerator.training_type_plugin._local_rank = rank
checkpoint_path = os.path.join(args.checkpoint_folder, f'mp_rank_{rank:02d}', args.checkpoint_name)
model = MegatronGPTModel.load_from_checkpoint(checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
model._save_restore_connector = NLPSaveRestoreConnector()
model.save_to(args.nemo_file_path)
logging.info(f'NeMo model saved to: {args.nemo_file_path}')
def main(cfg) -> None:
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'\n{OmegaConf.to_yaml(cfg)}')
plugins = [
NLPDDPPlugin(
no_ddp_communication_hook=True,
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),
hysteresis=cfg.model.get('hysteresis', 2),
)
plugins.append(
PipelineMixedPrecisionPlugin(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)
# 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
# load existing or init new soft prompt GPT model
if cfg.model.get("restore_path", None):
model = MegatronGPTPromptLearningModel.restore_from(
cfg.model.restore_path,
cfg.model,
trainer=trainer,
save_restore_connector=NLPSaveRestoreConnector())
else:
model = MegatronGPTPromptLearningModel(cfg.model, trainer=trainer)
trainer.fit(model)
def main():
parser = ArgumentParser()
parser.add_argument("--model_file",
type=str,
required=True,
help="Path to source .nemo file")
parser.add_argument("--target_file",
type=str,
required=True,
help="Path to write target .nemo file")
parser.add_argument("--tensor_model_parallel_size",
type=int,
required=True,
help="TP size of source model")
parser.add_argument("--target_tensor_model_parallel_size",
type=int,
required=True,
help="TP size of target model")
parser.add_argument(
"--model_class",
type=str,
default=
"nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTModel",
help=
"NeMo model class. This script should support all NeMo megatron models that use Tensor Parallel",
)
parser.add_argument("--precision",
default=16,
help="PyTorch Lightning Trainer precision flag")
args = parser.parse_args()
precision = args.precision
if args.precision in ["32", "16"]:
precision = int(float(args.precision))
tp_size = args.tensor_model_parallel_size
tgt_tp_size = args.target_tensor_model_parallel_size
cls = model_utils.import_class_by_path(args.model_class)
trainer = Trainer(devices=1,
plugins=NLPDDPPlugin(),
accelerator="cpu",
precision=precision)
app_state = AppState()
app_state.data_parallel_rank = 0
app_state.pipeline_model_parallel_size = 1 # not supported yet in this script
app_state.tensor_model_parallel_size = tp_size
app_state.model_parallel_size = app_state.pipeline_model_parallel_size * app_state.tensor_model_parallel_size
if tp_size > 1:
partitions = []
for i in range(tp_size):
app_state.tensor_model_parallel_rank = i
model = cls.restore_from(restore_path=args.model_file,
trainer=trainer,
map_location=torch.device("cpu"))
params = [p for _, p in model.named_parameters()]
partitions.append(params)
# app_state is being updated incorrectly during restore
app_state.data_parallel_rank = 0
app_state.pipeline_model_parallel_size = 1 # not supported yet in this script
app_state.tensor_model_parallel_size = tp_size
app_state.model_parallel_size = (
app_state.pipeline_model_parallel_size *
app_state.tensor_model_parallel_size)
model.cfg.tensor_model_parallel_size = 1
app_state.model_parallel_size = 1
trainer = Trainer(devices=1,
plugins=NLPDDPPlugin(),
accelerator="cpu",
precision=precision)
model = cls(model.cfg, trainer).to('cpu')
model._save_restore_connector = NLPSaveRestoreConnector()
if tgt_tp_size > 1:
merge_partition(model, partitions)
else:
merge_partition(model, partitions, args.target_file)
else:
app_state.model_parallel_size = 1
model = cls.restore_from(restore_path=args.model_file, trainer=trainer)
if tgt_tp_size > 1:
partitions = []
params = [p for _, p in model.named_parameters()]
partitions.append(params)
model.cfg.tensor_model_parallel_size = tgt_tp_size
app_state.model_parallel_size = tgt_tp_size
trainer = Trainer(devices=1,
plugins=NLPDDPPlugin(),
accelerator="cpu",
precision=precision)
model = cls(model.cfg, trainer).to('cpu')
model._save_restore_connector = NLPSaveRestoreConnector()
split_partition(model, partitions, tgt_tp_size, args.target_file)
logging.info("Successfully finished changing partitions!")
def convert(local_rank, rank, world_size, args):
app_state = AppState()
app_state.data_parallel_rank = 0
tensor_model_parallel_size = args.tensor_model_parallel_size
num_nodes = world_size // args.gpus_per_node
pipeline_model_parallel_size = world_size // args.tensor_model_parallel_size
assert args.pipeline_model_parallel_size == pipeline_model_parallel_size
trainer = Trainer(devices=args.gpus_per_node,
accelerator='gpu',
num_nodes=num_nodes)
app_state.pipeline_model_parallel_size = args.pipeline_model_parallel_size
app_state.tensor_model_parallel_size = args.tensor_model_parallel_size
app_state.model_parallel_size = app_state.tensor_model_parallel_size * app_state.pipeline_model_parallel_size
parallel_state.initialize_model_parallel(
tensor_model_parallel_size_=app_state.tensor_model_parallel_size,
pipeline_model_parallel_size_=app_state.pipeline_model_parallel_size,
)
app_state.pipeline_model_parallel_rank = parallel_state.get_pipeline_model_parallel_rank(
)
app_state.tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank(
)
pipeline_rank = rank // tensor_model_parallel_size
tensor_rank = app_state.tensor_model_parallel_rank
assert pipeline_rank == app_state.pipeline_model_parallel_rank
if tensor_model_parallel_size is not None and tensor_model_parallel_size > 1 and pipeline_model_parallel_size == 1:
# inject model parallel rank
checkpoint_path = os.path.join(args.checkpoint_folder,
f'mp_rank_{tensor_rank:02d}',
args.checkpoint_name)
elif tensor_model_parallel_size is not None and pipeline_model_parallel_size > 1:
checkpoint_path = os.path.join(
args.checkpoint_folder,
f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}',
args.checkpoint_name)
else:
checkpoint_path = os.path.join(args.checkpoint_folder,
args.checkpoint_name)
logging.info(f"loading checkpoint {checkpoint_path}")
if args.model_type == 'gpt':
## this dictionary is used to rename the model parameters
name_translate = {}
name_translate['transformer'] = 'encoder'
name_translate['.attention.'] = '.self_attention.'
# nemo megatron doesn't have _for_head key
name_translate['word_embeddings_for_head'] = 'word_embeddings'
checkpoint, consumed, steps, version = load_from_checkpoint(
MegatronGPTModel,
checkpoint_path,
hparams_file=args.hparams_file,
trainer=trainer,
translator=name_translate,
strict=False,
)
elif args.model_type == 'bert':
## this dictionary is used to rename the model parameters
name_translate = {}
name_translate['transformer'] = 'encoder'
name_translate['.attention.'] = '.self_attention.'
# nemo megatron doesn't have _for_head key
name_translate['word_embeddings_for_head'] = 'word_embeddings'
checkpoint, consumed, steps, version = load_from_checkpoint(
MegatronBertModel,
checkpoint_path,
hparams_file=args.hparams_file,
trainer=trainer,
translator=name_translate,
strict=False,
)
else:
raise NotImplemented("{} is not supported".format(args.model_type))
if torch.distributed.is_initialized():
torch.distributed.barrier()
if args.output_ckpt_file_path:
filepath = args.output_ckpt_file_path
base_dir = pathlib.Path(filepath).parent
filename_str = pathlib.Path(filepath).name
suffix = '.ckpt'
content = {}
if consumed is not None:
content['consumed'] = consumed
else:
content['consumed'] = 0
if steps is not None:
content['steps'] = steps
else:
content['steps'] = 0
filename = filename_str.format(**content) + suffix
checkpoint_path_output = inject_model_parallel_rank(
os.path.join(base_dir, filename))
trainer.accelerator.training_type_plugin.checkpoint_io.save_checkpoint(
checkpoint, checkpoint_path_output)
logging.info(
f'NeMo model checkpoint files saved to: {args.output_ckpt_file_path}'
)
if args.nemo_file_path:
if args.model_type == 'gpt':
model = load_model(MegatronGPTModel,
checkpoint,
strict=False,
trainer=trainer)
elif args.model_type == 'bert':
model = load_model(MegatronBertModel,
checkpoint,
strict=False,
trainer=trainer)
else:
raise NotImplemented("{} is not supported".format(args.model_type))
# verify tensor parallel rank id and pipeline parallel rank id matches
assert app_state.data_parallel_size == 1
assert app_state.tensor_model_parallel_size == tensor_model_parallel_size
assert app_state.tensor_model_parallel_rank == tensor_rank
assert app_state.pipeline_model_parallel_size == pipeline_model_parallel_size
assert app_state.pipeline_model_parallel_rank == pipeline_rank
model._save_restore_connector = NLPSaveRestoreConnector()
model.save_to(args.nemo_file_path)
logging.info(f'NeMo model saved to: {args.nemo_file_path}')
def nemo_convert(argv):
args = get_args(argv)
loglevel = logging.INFO
# assuming loglevel is bound to the string value obtained from the
# command line argument. Convert to upper case to allow the user to
# specify --log=DEBUG or --log=debug
if args.verbose is not None:
numeric_level = getattr(logging, args.verbose.upper(), None)
if not isinstance(numeric_level, int):
raise ValueError('Invalid log level: %s' % numeric_level)
loglevel = numeric_level
logger = logging.getLogger(__name__)
if logger.handlers:
for handler in logger.handlers:
logger.removeHandler(handler)
logging.basicConfig(level=loglevel,
format='%(asctime)s [%(levelname)s] %(message)s')
logging.info("Logging level set to {}".format(loglevel))
"""Convert a .nemo saved model trained on previous versions of nemo into a nemo fie with current version."""
nemo_in = args.source
out = args.out
# Create a PL trainer object which is required for restoring Megatron models
cfg_trainer = TrainerConfig(
gpus=1,
accelerator="ddp",
num_nodes=1,
# Need to set the following two to False as ExpManager will take care of them differently.
logger=False,
checkpoint_callback=False,
)
trainer = pl.Trainer(cfg_trainer)
logging.info("Restoring NeMo model from '{}'".format(nemo_in))
try:
# If the megatron based NLP model was trained on NeMo < 1.5, then we need to update the lm_checkpoint on the model config
if args.megatron_legacy:
if args.megatron_checkpoint:
connector = NLPSaveRestoreConnector()
model_cfg = ModelPT.restore_from(
restore_path=nemo_in,
save_restore_connector=connector,
trainer=trainer,
return_config=True)
OmegaConf.set_struct(model_cfg, True)
with open_dict(model_cfg):
model_cfg.language_model.lm_checkpoint = args.megatron_checkpoint
model_cfg['megatron_legacy'] = True
model_cfg['masked_softmax_fusion'] = False
model_cfg['bias_gelu_fusion'] = False
model = ModelPT.restore_from(
restore_path=nemo_in,
save_restore_connector=connector,
trainer=trainer,
override_config_path=model_cfg,
)
else:
logging.error(
"Megatron Checkpoint must be provided if Megatron legacy is chosen"
)
else:
model = ModelPT.restore_from(restore_path=nemo_in, trainer=trainer)
logging.info("Model {} restored from '{}'".format(
model.cfg.target, nemo_in))
# Save the model
model.save_to(out)
logging.info("Successfully converted to {}".format(out))
del model
except Exception as e:
logging.error(
"Failed to restore model from NeMo file : {}. Please make sure you have the latest NeMo package installed with [all] dependencies."
.format(nemo_in))
raise e
def __init__(self, cfg: DictConfig, trainer: Trainer = None, no_lm_init=False):
self.hidden_size = None
self.bert_model = None
vocab_file = None
nemo_file = None
config_dict = None
config_file = None
# tokenizer needs to get initialized before the super.__init__()
# as dataloaders and datasets need it to process the data
pretrain_model_name = ''
if cfg.get('language_model') and cfg.language_model.get('pretrained_model_name', ''):
pretrain_model_name = cfg.language_model.get('pretrained_model_name', '')
all_pretrained_megatron_bert_models = get_megatron_pretrained_bert_models()
if cfg.get('tokenizer'):
# Some models have their own tokenizer setup
if (
not hasattr(self, 'tokenizer')
and cfg.tokenizer.get('tokenizer_name')
and pretrain_model_name not in all_pretrained_megatron_bert_models
):
self.setup_tokenizer(cfg.tokenizer)
elif pretrain_model_name in all_pretrained_megatron_bert_models:
copy_cfg = copy.deepcopy(cfg)
bert_model = get_lm_model(
config_file=config_file,
config_dict=config_dict,
vocab_file=vocab_file,
trainer=trainer,
cfg=copy_cfg,
)
# set the tokenizer if it is not initialized explicitly
if (
(hasattr(self, 'tokenizer') and self.tokenizer is None) or not hasattr(self, 'tokenizer')
) and hasattr(bert_model, 'tokenizer'):
self.tokenizer = bert_model.tokenizer
if (
cfg.get('tokenizer')
and hasattr(cfg.get('tokenizer'), 'vocab_file')
and cfg.get('tokenizer').get('vocab_file')
):
vocab_file = self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file)
super().__init__(cfg, trainer)
# handles model parallel save and restore logic
self._save_restore_connector = NLPSaveRestoreConnector()
if cfg.get('language_model') and not no_lm_init:
if cfg.get('language_model').get('nemo_file'):
nemo_file = self.register_artifact('language_model.nemo_file', cfg.language_model.nemo_file)
if cfg.get('language_model').get('config'):
config_dict = OmegaConf.to_container(cfg.language_model.config)
if cfg.get('language_model').get('config_file'):
config_file = self.register_artifact('language_model.config_file', cfg.language_model.config_file)
bert_model = get_lm_model(
config_file=config_file, config_dict=config_dict, vocab_file=vocab_file, trainer=trainer, cfg=cfg,
)
# set the tokenizer if it is not initialized explicitly
if ((hasattr(self, 'tokenizer') and self.tokenizer is None) or not hasattr(self, 'tokenizer')) and hasattr(
bert_model, 'tokenizer'
):
self.tokenizer = bert_model.tokenizer
# Required to pull up the config for MegatronBert models
self.pretrained_model_name = cfg.language_model.pretrained_model_name
# register encoder config
self.register_bert_model()
if (
cfg.tokenizer is not None
and cfg.tokenizer.get("tokenizer_name", "") is not None
and "megatron" in cfg.tokenizer.get("tokenizer_name", "")
) or pretrain_model_name in all_pretrained_megatron_bert_models:
self.hidden_size = bert_model.cfg.hidden_size
else:
self.hidden_size = bert_model.config.hidden_size
if cfg.get('language_model') and not no_lm_init:
self.bert_model = bert_model
def __init__(self, cfg: DictConfig, trainer: Trainer):
super().__init__(cfg, trainer)
self.cfg = cfg
# Load pretrained GPT model and tokenizer
if cfg.get('language_model_path', None):
self.frozen_model = MegatronGPTModel.restore_from(
cfg.get('language_model_path'),
trainer=trainer,
save_restore_connector=NLPSaveRestoreConnector(),
)
# Freeze all GPT model weights for prompt-tuning/p-tuning
self.frozen_model.freeze()
self.tokenizer = self.frozen_model.tokenizer
self.float_type = self.frozen_model.model.language_model.encoder.layers[
0].dtype
self.hidden_size = self.frozen_model.cfg.hidden_size
self.word_embeddings = self.frozen_model.model.language_model.embedding.word_embeddings
self.existing_tasks = list(self.cfg.get('existing_tasks', []))
self.new_tasks = list(self.cfg.get('new_tasks', []))
# Load templates for assigning virtual prompt token positions
self.load_task_templates(self.cfg.task_templates)
# Prompt table stores all task embeddings, p-tuning virtual prompts get added to the table after training
self.prompt_table = PromptTable(
existing_tasks=self.existing_tasks,
task_templates=self.task_templates,
task_id_num_to_name=self.task_id_num_to_name,
hidden_size=self.hidden_size,
)
self._prompt_table_key = VirtualPromptSource.PROMPT_TABLE.value
self._prompt_encoder_key = VirtualPromptSource.PROMPT_ENCODER.value
# Prepare pseudo token ids for virtual/virtual prompt tokens
self.pseudo_tokens = get_pseudo_tokens(self.max_virtual_tokens)
self.tokenizer.add_special_tokens(
{'additional_special_tokens': self.pseudo_tokens})
self.pseudo_token_ids = self.tokenizer.tokens_to_ids(
self.pseudo_tokens)
self.pseudo_token_ids_start = self.pseudo_token_ids[0]
self.pad_token_id = self.tokenizer.pad_id if self.tokenizer.pad_id is not None else self.tokenizer.unk_id
self.virtual_prompt_style = VirtualPromptStyle(
cfg.virtual_prompt_style)
# Prompt tuning stores virtual prompts in the prompt table and tunes their weight directly
if self.virtual_prompt_style in [
VirtualPromptStyle.PROMPT_TUNING, VirtualPromptStyle.INFERENCE
]:
self.virtual_prompt_source = VirtualPromptSource.PROMPT_TABLE
# P-Tuning uses an LSTM Encoder to produce virtual token embeddings
elif self.virtual_prompt_style == VirtualPromptStyle.P_TUNING:
self.virtual_prompt_source = VirtualPromptSource.PROMPT_ENCODER
else:
raise ValueError(
f"\nvirtual prompt style '{cfg.virtual_prompt_style}' not recognized, please use one of 'prompt-tuning' or 'p-tuning'"
)
self._reduced_loss_buffer = []
self._inference_config = None
if self.trainer.precision == 32:
self.autocast_dtype = torch.float
elif self.trainer.precision == 16:
self.autocast_dtype = torch.half
elif self.trainer.precision == 'bf16':
self.autocast_dtype = torch.bfloat16
else:
raise ValueError('precision must be in [32, 16, "bf16"]')
# make sure the default pytorch lightning gradient clipping in the basemodel
self.grad_clip_pl_default = True
# no support of amp o2
self.megatron_amp_o2 = False
def __init__(self, cfg: DictConfig, trainer: Trainer):
super().__init__(cfg, trainer)
self.cfg = cfg
# Load pretrained GPT model and tokenizer
self.model = MegatronGPTModel.restore_from(
self.register_artifact('language_model_path',
cfg.get('language_model_path', None)),
trainer=trainer,
save_restore_connector=NLPSaveRestoreConnector(),
)
# Freeze all GPT model weights for prompt-tuning/p-tuning
if not cfg.lm_finetune:
self.model.freeze()
self.tokenizer = self.model.tokenizer
self.float_type = self.model.model.language_model.encoder.layers[
0].dtype
self.hidden_size = self.model.cfg.hidden_size
self.word_embeddings = self.model.model.language_model.embedding.word_embeddings
self.existing_tasks = list(self.cfg.get('existing_tasks', []))
self.new_tasks = list(self.cfg.get('new_tasks', []))
# Load templates for assigning virtual prompt token positions
self.load_task_templates(self.cfg.task_templates)
# Prompt table stores all task embeddings, p-tuning virtual prompts get added to the table after training
self.prompt_table = PromptTable(
existing_tasks=self.existing_tasks,
task_templates=self.task_templates,
task_id_num_to_name=self.task_id_num_to_name,
hidden_size=self.hidden_size,
)
# Prepare pseudo token ids for virtual/virtual prompt tokens
self.pseudo_token_base = cfg.pseudo_token_base
self.pseudo_tokens = [
self.pseudo_token_base + str(i)
for i in range(self.max_virtual_tokens)
]
self.tokenizer.add_special_tokens(
{'additional_special_tokens': self.pseudo_tokens})
self.pseudo_token_ids = self.tokenizer.tokens_to_ids(
self.pseudo_tokens)
self.pseudo_token_ids_start = self.pseudo_token_ids[0]
self.pad_token_id = self.tokenizer.pad_id if self.tokenizer.pad_id is not None else self.tokenizer.unk_id
self.virtual_prompt_style = cfg.virtual_prompt_style.lower()
# Prompt tuning stores virtual prompts in the prompt table and tunes their weight directly
if self.virtual_prompt_style in ['prompt-tuning', 'inference']:
self.virtual_prompt_source = 'prompt-table'
# P-Tuning uses an LSTM Encoder to produce virtual token embeddings
elif self.virtual_prompt_style == 'p-tuning':
self.virtual_prompt_source = 'prompt-encoder'
else:
raise ValueError(
f"\nvirtual prompt style '{cfg.virtual_prompt_type}' not recognized, please use one of 'prompt-tuning' or 'p-tuning'"
)
self._reduced_loss_buffer = []
self._inference_config = None
if self.trainer.precision == 32:
self.autocast_dtype = torch.float
elif self.trainer.precision == 16:
self.autocast_dtype = torch.half
elif self.trainer.precision == 'bf16':
self.autocast_dtype = torch.bfloat16
else:
raise ValueError('precision must be in [32, 16, "bf16"]')
def main():
parser = ArgumentParser()
parser.add_argument("--model_file",
type=str,
default="",
required=True,
help="Pass path to model's .nemo file")
parser.add_argument("--prompt",
type=str,
default="",
required=True,
help="Prompt for the model (a text to complete)")
parser.add_argument("--tokens_to_generate",
type=int,
default="16",
required=False,
help="How many tokens to add to prompt")
parser.add_argument(
"--tensor_model_parallel_size",
type=int,
default=1,
required=False,
)
parser.add_argument(
"--pipeline_model_parallel_size",
type=int,
default=1,
required=False,
)
parser.add_argument(
"--pipeline_model_parallel_split_rank",
type=int,
default=0,
required=False,
)
parser.add_argument("--precision",
default="16",
type=str,
help="PyTorch Lightning Trainer precision flag")
args = parser.parse_args()
# cast precision to int if 32 or 16
if args.precision in ["32", "16"]:
args.precision = int(float(args.precision))
# trainer required for restoring model parallel models
trainer = Trainer(
plugins=NLPDDPPlugin(),
devices=args.tensor_model_parallel_size *
args.pipeline_model_parallel_size,
accelerator='gpu',
precision=args.precision,
)
app_state = AppState()
if args.tensor_model_parallel_size > 1 or args.pipeline_model_parallel_size > 1:
app_state.model_parallel_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size
(
app_state.tensor_model_parallel_rank,
app_state.pipeline_model_parallel_rank,
app_state.model_parallel_size,
app_state.data_parallel_size,
app_state.pipeline_model_parallel_split_rank,
) = fake_initialize_model_parallel(
world_size=app_state.model_parallel_size,
rank=trainer.global_rank,
tensor_model_parallel_size_=args.tensor_model_parallel_size,
pipeline_model_parallel_size_=args.pipeline_model_parallel_size,
pipeline_model_parallel_split_rank_=args.
pipeline_model_parallel_split_rank,
)
model = MegatronT5Model.restore_from(
restore_path=args.model_file,
trainer=trainer,
save_restore_connector=NLPSaveRestoreConnector(),
)
model.freeze()
request = {
"prompt": args.prompt,
"tokens_to_generate": args.tokens_to_generate,
}
dataset = T5RequestDataset(request, model.tokenizer)
request_dl = DataLoader(dataset)
response = trainer.predict(model, request_dl)
print("***************************")
print(response)
print("***************************")
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=True,
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(PipelineMixedPrecisionPlugin(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
if cfg.model.resume_from_checkpoint is not None:
resume_from_checkpoint = cfg.model.resume_from_checkpoint
else:
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,)
# Get the T5 Base configuration.
t5_cfg = MegatronT5FinetuneModel.restore_from(
restore_path=cfg.model.restore_from_path, trainer=trainer, return_config=True
)
# Override the T5 configuration with the one from the config file.
OmegaConf.set_struct(t5_cfg, True)
with open_dict(t5_cfg):
t5_cfg.masked_softmax_fusion = False
t5_cfg.megatron_amp_O2 = cfg.model.get('megatron_amp_O2', False)
t5_cfg.hidden_dropout = cfg.model.get('hidden_dropout', 0.1)
t5_cfg.attention_dropout = cfg.model.get('attention_dropout', 0.1)
t5_cfg.data = cfg.model.data
t5_cfg.precision = cfg.trainer.precision
t5_cfg.optim = cfg.model.optim
t5_cfg.micro_batch_size = cfg.model.data.train_ds.micro_batch_size
t5_cfg.global_batch_size = cfg.model.data.train_ds.global_batch_size
# XNLI has eval languages in the yaml config.
if hasattr(cfg.model, 'eval_languages'):
t5_cfg.eval_languages = cfg.model.eval_languages
if hasattr(cfg.model.data.train_ds, 'task_name'):
model = MegatronT5GLUEModel.restore_from(
restore_path=cfg.model.restore_from_path,
trainer=trainer,
override_config_path=t5_cfg,
save_restore_connector=NLPSaveRestoreConnector(),
)
else:
model = MegatronT5FinetuneModel.restore_from(
restore_path=cfg.model.restore_from_path,
trainer=trainer,
override_config_path=t5_cfg,
save_restore_connector=NLPSaveRestoreConnector(),
)
trainer.fit(model)
trainer.validate(model)
if hasattr(cfg.model.data, 'test_ds'):
trainer.test(model)
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=True,
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(
PipelineMixedPrecisionPlugin(precision=cfg.trainer.precision,
device='cuda',
scaler=scaler))
if cfg.get('cluster_type', None) == 'BCP':
plugins.append(TorchElasticEnvironment())
trainer = Trainer(plugins=plugins,
**cfg.trainer,
callbacks=[ModelSummary(max_depth=3)])
# tokenizers will be trained and and tarred training data will be created if needed
# model config is then updated
if cfg.model.preproc_out_dir is not None:
MTDataPreproc(cfg=cfg.model, trainer=trainer)
exp_manager(trainer, cfg.exp_manager)
# update resume from checkpoint found by exp_manager
if cfg.model.resume_from_checkpoint is not None:
resume_from_checkpoint = cfg.model.resume_from_checkpoint
else:
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
if hasattr(cfg.model, 'pretrained_model_path'
) and cfg.model.pretrained_model_path is not None:
if not hasattr(cfg.model, 'pretrained_model_type'):
raise ValueError(f"Pretrained model type must be in [T5, BART].")
assert cfg.model.pretrained_model_type in ['T5', 'BART']
if cfg.model.pretrained_model_type == 'T5':
pretrained_cfg = MegatronT5Model.restore_from(
cfg.model.pretrained_model_path,
trainer=trainer,
return_config=True)
else:
pretrained_cfg = MegatronBARTModel.restore_from(
cfg.model.pretrained_model_path,
trainer=trainer,
return_config=True)
OmegaConf.set_struct(pretrained_cfg, True)
with open_dict(pretrained_cfg):
pretrained_cfg.masked_softmax_fusion = False
# Set source and target language/multilingual
pretrained_cfg.src_language = cfg.model.src_language
pretrained_cfg.tgt_language = cfg.model.tgt_language
pretrained_cfg.multilingual = cfg.model.multilingual
pretrained_cfg.shared_tokenizer = True
# Max generation delta
pretrained_cfg.max_generation_delta = cfg.model.max_generation_delta
# Set label smoothing
pretrained_cfg.label_smoothing = cfg.model.label_smoothing
# Set tokenizer paths:
pretrained_cfg.encoder_tokenizer = pretrained_cfg.tokenizer
pretrained_cfg.decoder_tokenizer = pretrained_cfg.tokenizer
# Pre-trained models should use the legacy sentencepiece tokenizer ex: mT5
pretrained_cfg.encoder_tokenizer.sentencepiece_legacy = True
pretrained_cfg.decoder_tokenizer.sentencepiece_legacy = True
# Override dropout
pretrained_cfg.hidden_dropout = cfg.model.hidden_dropout
pretrained_cfg.attention_dropout = cfg.model.attention_dropout
# Override precision
pretrained_cfg.precision = cfg.model.precision # Set above from trainer.precision
# Override data and global/micro batch size.
pretrained_cfg.train_ds = cfg.model.train_ds
pretrained_cfg.validation_ds = cfg.model.validation_ds
pretrained_cfg.test_ds = cfg.model.test_ds
pretrained_cfg.micro_batch_size = cfg.model.micro_batch_size
pretrained_cfg.global_batch_size = cfg.model.global_batch_size
# Class target for the new class being restored.
pretrained_cfg.target = (
"nemo.collections.nlp.models.machine_translation.megatron_nmt_model.MegatronNMTModel"
)
# Optimizer overrides.
pretrained_cfg.optim = cfg.model.optim
model = MegatronNMTModel.restore_from(
cfg.model.pretrained_model_path,
trainer=trainer,
override_config_path=pretrained_cfg,
save_restore_connector=NLPSaveRestoreConnector(),
)
else:
model = MegatronNMTModel(cfg.model, trainer)
if cfg.do_training:
trainer.fit(model)
if cfg.do_testing:
trainer.test(model)
def __init__(self, cfg: DictConfig, trainer: Trainer = None):
super().__init__(cfg, trainer)
# handles model parallel save and restore logic
self._save_restore_connector = NLPSaveRestoreConnector()
self.set_world_size(trainer)