def setup_test_data(self, test_data_config: Optional[DictConfig]):
self._test_dl = self._setup_eval_dataloader_from_config(cfg=test_data_config)
# instantiate Torchmetric for each test dataloader
if self._test_dl is not None:
for dataloader_idx in range(len(self._test_dl)):
if dataloader_idx == 0:
setattr(
self, f'test_loss', GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True),
)
else:
setattr(
self,
f'test_loss_{dataloader_idx}',
GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True),
)
def __init__(self, cfg: DictConfig, trainer: Trainer = None):
# Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
self.world_size = 1
if trainer is not None:
self.world_size = trainer.num_nodes * trainer.num_gpus
cfg = model_utils.convert_model_config_to_dict_config(cfg)
cfg = model_utils.maybe_update_config_version(cfg)
# Instantiates tokenizer and register to be saved with NeMo Model archive
# After this call, ther will be self.tokenizer which can convert between tokens and token_ids.
self.setup_tokenizer(
tokenizer_name=cfg.tokenizer.get("tokenizer_name", "yttm"),
tokenizer_model=cfg.tokenizer.get("tokenizer_model", None),
vocab_file=cfg.tokenizer.get("vocab_file", None),
bpe_dropout=cfg.tokenizer.get("bpe_dropout", 0.0),
special_tokens=cfg.tokenizer.get("special_tokens", {}))
# init superclass
super().__init__(cfg=cfg, trainer=trainer)
# make vocabulary size divisible by 8 for fast fp16 training
vocab_size = 8 * math.ceil(self.tokenizer.vocab_size / 8)
# encoder from NeMo, Megatron-LM, or HuggingFace
encoder_cfg_dict = OmegaConf.to_container(cfg.get('encoder'))
encoder_cfg_dict['vocab_size'] = vocab_size
library = encoder_cfg_dict.pop('library', 'nemo')
model_name = encoder_cfg_dict.pop('model_name', None)
pretrained = encoder_cfg_dict.pop('pretrained', False)
self.encoder = get_transformer(
library=library,
model_name=model_name,
pretrained=pretrained,
config_dict=encoder_cfg_dict,
encoder=True,
pre_ln_final_layer_norm=encoder_cfg_dict.get(
'pre_ln_final_layer_norm',
encoder_cfg_dict.get('pre_ln', True)),
)
self.log_softmax = TokenClassifier(
hidden_size=self.encoder.hidden_size,
num_classes=vocab_size,
activation=cfg.head.activation,
log_softmax=cfg.head.log_softmax,
dropout=cfg.head.dropout,
use_transformer_init=cfg.head.use_transformer_init,
)
# tie weights of embedding and softmax matrices
self.log_softmax.mlp.layer0.weight = self.encoder.embedding.token_embedding.weight
std_init_range = 1 / self.encoder.hidden_size**0.5
# initialize weights if not using pretrained encoder
if not self._cfg.encoder.get('pretrained', False):
self.encoder.apply(lambda module: transformer_weights_init(
module, std_init_range))
self.log_softmax.apply(
lambda module: transformer_weights_init(module, std_init_range))
self.loss_fn = SmoothedCrossEntropyLoss(
pad_id=self.tokenizer.pad_id, label_smoothing=cfg.label_smoothing)
self.eval_loss_fn = SmoothedCrossEntropyLoss(
pad_id=self.tokenizer.pad_id)
self.eval_loss = GlobalAverageLossMetric(dist_sync_on_step=False,
take_avg_loss=True)
self.eval_ppl = SequencePerplexity()
def __init__(self, cfg: MTEncDecModelConfig, trainer: Trainer = None):
cfg = model_utils.convert_model_config_to_dict_config(cfg)
# Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
# Global_rank and local_rank is set by LightningModule in Lightning 1.2.0
self.world_size = 1
if trainer is not None:
self.world_size = trainer.num_nodes * trainer.num_gpus
cfg = model_utils.maybe_update_config_version(cfg)
self.src_language: str = cfg.get("src_language", None)
self.tgt_language: str = cfg.get("tgt_language", None)
# Instantiates tokenizers and register to be saved with NeMo Model archive
# After this call, ther will be self.encoder_tokenizer and self.decoder_tokenizer
# Which can convert between tokens and token_ids for SRC and TGT languages correspondingly.
self.setup_enc_dec_tokenizers(
encoder_tokenizer_name=cfg.encoder_tokenizer.tokenizer_name,
encoder_tokenizer_model=cfg.encoder_tokenizer.tokenizer_model,
encoder_bpe_dropout=cfg.encoder_tokenizer.get('bpe_dropout', 0.0),
decoder_tokenizer_name=cfg.decoder_tokenizer.tokenizer_name,
decoder_tokenizer_model=cfg.decoder_tokenizer.tokenizer_model,
decoder_bpe_dropout=cfg.decoder_tokenizer.get('bpe_dropout', 0.0),
)
# After this call, the model will have self.source_processor and self.target_processor objects
self.setup_pre_and_post_processing_utils(source_lang=self.src_language,
target_lang=self.tgt_language)
# TODO: Why is this base constructor call so late in the game?
super().__init__(cfg=cfg, trainer=trainer)
# TODO: use get_encoder function with support for HF and Megatron
self.encoder = TransformerEncoderNM(
vocab_size=self.encoder_vocab_size,
hidden_size=cfg.encoder.hidden_size,
num_layers=cfg.encoder.num_layers,
inner_size=cfg.encoder.inner_size,
max_sequence_length=cfg.encoder.max_sequence_length if hasattr(
cfg.encoder, 'max_sequence_length') else 512,
embedding_dropout=cfg.encoder.embedding_dropout if hasattr(
cfg.encoder, 'embedding_dropout') else 0.0,
learn_positional_encodings=cfg.encoder.learn_positional_encodings
if hasattr(cfg.encoder, 'learn_positional_encodings') else False,
num_attention_heads=cfg.encoder.num_attention_heads,
ffn_dropout=cfg.encoder.ffn_dropout,
attn_score_dropout=cfg.encoder.attn_score_dropout,
attn_layer_dropout=cfg.encoder.attn_layer_dropout,
hidden_act=cfg.encoder.hidden_act,
mask_future=cfg.encoder.mask_future,
pre_ln=cfg.encoder.pre_ln,
)
# TODO: user get_decoder function with support for HF and Megatron
self.decoder = TransformerDecoderNM(
vocab_size=self.decoder_vocab_size,
hidden_size=cfg.decoder.hidden_size,
num_layers=cfg.decoder.num_layers,
inner_size=cfg.decoder.inner_size,
max_sequence_length=cfg.decoder.max_sequence_length if hasattr(
cfg.decoder, 'max_sequence_length') else 512,
embedding_dropout=cfg.decoder.embedding_dropout if hasattr(
cfg.decoder, 'embedding_dropout') else 0.0,
learn_positional_encodings=cfg.decoder.learn_positional_encodings
if hasattr(cfg.decoder, 'learn_positional_encodings') else False,
num_attention_heads=cfg.decoder.num_attention_heads,
ffn_dropout=cfg.decoder.ffn_dropout,
attn_score_dropout=cfg.decoder.attn_score_dropout,
attn_layer_dropout=cfg.decoder.attn_layer_dropout,
hidden_act=cfg.decoder.hidden_act,
pre_ln=cfg.decoder.pre_ln,
)
self.log_softmax = TokenClassifier(
hidden_size=self.decoder.hidden_size,
num_classes=self.decoder_vocab_size,
activation=cfg.head.activation,
log_softmax=cfg.head.log_softmax,
dropout=cfg.head.dropout,
use_transformer_init=cfg.head.use_transformer_init,
)
self.beam_search = BeamSearchSequenceGenerator(
embedding=self.decoder.embedding,
decoder=self.decoder.decoder,
log_softmax=self.log_softmax,
max_sequence_length=self.decoder.max_sequence_length,
beam_size=cfg.beam_size,
bos=self.decoder_tokenizer.bos_id,
pad=self.decoder_tokenizer.pad_id,
eos=self.decoder_tokenizer.eos_id,
len_pen=cfg.len_pen,
max_delta_length=cfg.max_generation_delta,
)
# tie weights of embedding and softmax matrices
self.log_softmax.mlp.layer0.weight = self.decoder.embedding.token_embedding.weight
# TODO: encoder and decoder with different hidden size?
std_init_range = 1 / self.encoder.hidden_size**0.5
self.apply(
lambda module: transformer_weights_init(module, std_init_range))
self.loss_fn = SmoothedCrossEntropyLoss(
pad_id=self.decoder_tokenizer.pad_id,
label_smoothing=cfg.label_smoothing)
self.eval_loss = GlobalAverageLossMetric(dist_sync_on_step=False,
take_avg_loss=True)
def _loss_class_test(
rank: int,
worldsize: int,
loss_sum_or_avg: Optional[torch.Tensor],
num_measurements: Optional[torch.Tensor],
dist_sync_on_step: bool,
take_avg_loss: bool,
check_dist_sync_on_step: bool = True,
check_batch: bool = True,
atol: float = 1e-8,
):
""" Utility function doing the actual comparison between lightning class metric
and reference metric.
Args:
rank: rank of current process
worldsize: number of processes
loss_sum_or_avg: a one dimensional float torch tensor with loss sums or means.
num_measurements: a one dimensional integer torch tensor with number of values on which sums or means from
``loss_sum_or_avg`` were computed.
dist_sync_on_step: bool, if true will synchronize metric state across processes at each call of the
method :meth:`forward()`
take_avg_loss: dict with additional arguments used for class initialization
check_dist_sync_on_step: bool, if true will check if the metric is also correctly
calculated per batch per device (and not just at the end)
check_batch: bool, if true will check if the metric is also correctly
calculated across devices for each batch (and not just at the end)
"""
# Instantiate lightning metric
loss_metric = GlobalAverageLossMetric(compute_on_step=True,
dist_sync_on_step=dist_sync_on_step,
take_avg_loss=take_avg_loss)
# verify loss works after being loaded from pickled state
pickled_metric = pickle.dumps(loss_metric)
loss_metric = pickle.loads(pickled_metric)
for i in range(rank, NUM_BATCHES, worldsize):
batch_result = loss_metric(loss_sum_or_avg[i], num_measurements[i])
if loss_metric.dist_sync_on_step:
if rank == 0:
ddp_loss_sum_or_avg = torch.stack(
[loss_sum_or_avg[i + r] for r in range(worldsize)])
ddp_num_measurements = torch.stack(
[num_measurements[i + r] for r in range(worldsize)])
sk_batch_result = reference_loss_func(ddp_loss_sum_or_avg,
ddp_num_measurements,
take_avg_loss)
# assert for dist_sync_on_step
if check_dist_sync_on_step:
if sk_batch_result.isnan():
assert batch_result.isnan()
else:
assert np.allclose(
batch_result.numpy(), sk_batch_result, atol=atol
), f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}, i = {i}"
else:
ls = loss_sum_or_avg[i:i + 1]
nm = num_measurements[i:i + 1]
sk_batch_result = reference_loss_func(ls, nm, take_avg_loss)
# assert for batch
if check_batch:
if sk_batch_result.isnan():
assert batch_result.isnan()
else:
assert np.allclose(
batch_result.numpy(), sk_batch_result, atol=atol
), f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}, i = {i}"
# check on all batches on all ranks
result = loss_metric.compute()
assert isinstance(result, torch.Tensor)
sk_result = reference_loss_func(loss_sum_or_avg, num_measurements,
take_avg_loss)
# assert after aggregation
if sk_result.isnan():
assert result.isnan()
else:
assert np.allclose(
result.numpy(), sk_result,
atol=atol), f"result = {result.numpy()}, sk_result = {sk_result}"
def __init__(self, cfg: MTEncDecModelConfig, trainer: Trainer = None):
cfg = model_utils.convert_model_config_to_dict_config(cfg)
# Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
self.global_rank = 0
self.world_size = 1
if trainer is not None:
self.global_rank = (trainer.node_rank *
trainer.num_gpus) + trainer.local_rank
self.world_size = trainer.num_nodes * trainer.num_gpus
cfg = model_utils.maybe_update_config_version(cfg)
self.setup_enc_dec_tokenizers(cfg)
super().__init__(cfg=cfg, trainer=trainer)
# TODO: use get_encoder function with support for HF and Megatron
self.encoder = TransformerEncoderNM(
vocab_size=self.encoder_vocab_size,
hidden_size=cfg.encoder.hidden_size,
num_layers=cfg.encoder.num_layers,
inner_size=cfg.encoder.inner_size,
max_sequence_length=cfg.encoder.max_sequence_length if hasattr(
cfg.encoder, 'max_sequence_length') else 512,
embedding_dropout=cfg.encoder.embedding_dropout if hasattr(
cfg.encoder, 'embedding_dropout') else 0.0,
learn_positional_encodings=cfg.encoder.learn_positional_encodings
if hasattr(cfg.encoder, 'learn_positional_encodings') else False,
num_attention_heads=cfg.encoder.num_attention_heads,
ffn_dropout=cfg.encoder.ffn_dropout,
attn_score_dropout=cfg.encoder.attn_score_dropout,
attn_layer_dropout=cfg.encoder.attn_layer_dropout,
hidden_act=cfg.encoder.hidden_act,
mask_future=cfg.encoder.mask_future,
pre_ln=cfg.encoder.pre_ln,
)
# TODO: user get_decoder function with support for HF and Megatron
self.decoder = TransformerDecoderNM(
vocab_size=self.decoder_vocab_size,
hidden_size=cfg.decoder.hidden_size,
num_layers=cfg.decoder.num_layers,
inner_size=cfg.decoder.inner_size,
max_sequence_length=cfg.decoder.max_sequence_length if hasattr(
cfg.decoder, 'max_sequence_length') else 512,
embedding_dropout=cfg.decoder.embedding_dropout if hasattr(
cfg.decoder, 'embedding_dropout') else 0.0,
learn_positional_encodings=cfg.decoder.learn_positional_encodings
if hasattr(cfg.decoder, 'learn_positional_encodings') else False,
num_attention_heads=cfg.decoder.num_attention_heads,
ffn_dropout=cfg.decoder.ffn_dropout,
attn_score_dropout=cfg.decoder.attn_score_dropout,
attn_layer_dropout=cfg.decoder.attn_layer_dropout,
hidden_act=cfg.decoder.hidden_act,
pre_ln=cfg.decoder.pre_ln,
)
self.log_softmax = TokenClassifier(
hidden_size=self.decoder.hidden_size,
num_classes=self.decoder_vocab_size,
activation=cfg.head.activation,
log_softmax=cfg.head.log_softmax,
dropout=cfg.head.dropout,
use_transformer_init=cfg.head.use_transformer_init,
)
self.beam_search = BeamSearchSequenceGenerator(
embedding=self.decoder.embedding,
decoder=self.decoder.decoder,
log_softmax=self.log_softmax,
max_sequence_length=self.decoder.max_sequence_length,
beam_size=cfg.beam_size,
bos=self.decoder_tokenizer.bos_id,
pad=self.decoder_tokenizer.pad_id,
eos=self.decoder_tokenizer.eos_id,
len_pen=cfg.len_pen,
max_delta_length=cfg.max_generation_delta,
)
# tie weights of embedding and softmax matrices
self.log_softmax.mlp.layer0.weight = self.decoder.embedding.token_embedding.weight
# TODO: encoder and decoder with different hidden size?
std_init_range = 1 / self.encoder.hidden_size**0.5
self.apply(
lambda module: transformer_weights_init(module, std_init_range))
self.loss_fn = SmoothedCrossEntropyLoss(
pad_id=self.decoder_tokenizer.pad_id,
label_smoothing=cfg.label_smoothing)
self.eval_loss = GlobalAverageLossMetric(dist_sync_on_step=False,
take_avg_loss=True)
def __init__(self, cfg: MTEncDecModelConfig, trainer: Trainer = None):
cfg = model_utils.convert_model_config_to_dict_config(cfg)
# Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
# Global_rank and local_rank is set by LightningModule in Lightning 1.2.0
self.world_size = 1
if trainer is not None:
self.world_size = trainer.num_nodes * trainer.num_gpus
cfg = model_utils.maybe_update_config_version(cfg)
self.src_language: str = cfg.get("src_language", None)
self.tgt_language: str = cfg.get("tgt_language", None)
# Instantiates tokenizers and register to be saved with NeMo Model archive
# After this call, ther will be self.encoder_tokenizer and self.decoder_tokenizer
# Which can convert between tokens and token_ids for SRC and TGT languages correspondingly.
self.setup_enc_dec_tokenizers(
encoder_tokenizer_library=cfg.encoder_tokenizer.get('library', 'yttm'),
encoder_tokenizer_model=cfg.encoder_tokenizer.get('tokenizer_model'),
encoder_bpe_dropout=cfg.encoder_tokenizer.get('bpe_dropout', 0.0),
encoder_model_name=cfg.encoder.get('model_name') if hasattr(cfg.encoder, 'model_name') else None,
decoder_tokenizer_library=cfg.decoder_tokenizer.get('library', 'yttm'),
decoder_tokenizer_model=cfg.decoder_tokenizer.tokenizer_model,
decoder_bpe_dropout=cfg.decoder_tokenizer.get('bpe_dropout', 0.0),
decoder_model_name=cfg.decoder.get('model_name') if hasattr(cfg.decoder, 'model_name') else None,
)
# After this call, the model will have self.source_processor and self.target_processor objects
self.setup_pre_and_post_processing_utils(source_lang=self.src_language, target_lang=self.tgt_language)
# TODO: Why is this base constructor call so late in the game?
super().__init__(cfg=cfg, trainer=trainer)
# encoder from NeMo, Megatron-LM, or HuggingFace
encoder_cfg_dict = OmegaConf.to_container(cfg.get('encoder'))
encoder_cfg_dict['vocab_size'] = self.encoder_vocab_size
library = encoder_cfg_dict.pop('library', 'nemo')
model_name = encoder_cfg_dict.pop('model_name', None)
pretrained = encoder_cfg_dict.pop('pretrained', False)
self.encoder = get_transformer(
library=library, model_name=model_name, pretrained=pretrained, config_dict=encoder_cfg_dict, encoder=True,
)
# decoder from NeMo, Megatron-LM, or HuggingFace
decoder_cfg_dict = OmegaConf.to_container(cfg.get('decoder'))
decoder_cfg_dict['vocab_size'] = self.decoder_vocab_size
library = decoder_cfg_dict.pop('library', 'nemo')
model_name = decoder_cfg_dict.pop('model_name', None)
pretrained = decoder_cfg_dict.pop('pretrained', False)
decoder_cfg_dict['hidden_size'] = self.encoder.hidden_size
self.decoder = get_transformer(
library=library, model_name=model_name, pretrained=pretrained, config_dict=decoder_cfg_dict, encoder=False,
)
self.log_softmax = TokenClassifier(
hidden_size=self.decoder.hidden_size,
num_classes=self.decoder_vocab_size,
activation=cfg.head.activation,
log_softmax=cfg.head.log_softmax,
dropout=cfg.head.dropout,
use_transformer_init=cfg.head.use_transformer_init,
)
self.beam_search = BeamSearchSequenceGenerator(
embedding=self.decoder.embedding,
decoder=self.decoder.decoder,
log_softmax=self.log_softmax,
max_sequence_length=self.decoder.max_sequence_length,
beam_size=cfg.beam_size,
bos=self.decoder_tokenizer.bos_id,
pad=self.decoder_tokenizer.pad_id,
eos=self.decoder_tokenizer.eos_id,
len_pen=cfg.len_pen,
max_delta_length=cfg.max_generation_delta,
)
# tie weights of embedding and softmax matrices
self.log_softmax.mlp.layer0.weight = self.decoder.embedding.token_embedding.weight
# TODO: encoder and decoder with different hidden size?
std_init_range = 1 / self.encoder.hidden_size ** 0.5
# initialize weights if not using pretrained encoder/decoder
if not self._cfg.encoder.get('pretrained', False):
self.encoder.apply(lambda module: transformer_weights_init(module, std_init_range))
if not self._cfg.decoder.get('pretrained', False):
self.decoder.apply(lambda module: transformer_weights_init(module, std_init_range))
self.log_softmax.apply(lambda module: transformer_weights_init(module, std_init_range))
self.loss_fn = SmoothedCrossEntropyLoss(
pad_id=self.decoder_tokenizer.pad_id, label_smoothing=cfg.label_smoothing
)
self.eval_loss = GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True)