def __init__(self,
pos_embed: DictConfig,
transformer: DictConfig,
layer_drop: float = 0.0):
super().__init__(**transformer) # see nlp.collections
# positional convolutional embeddings
emb_dim = pos_embed.embedding_dim
self.pos_conv = nn.Conv1d(
emb_dim,
emb_dim,
kernel_size=pos_embed.conv_pos,
padding=pos_embed.conv_pos //
2, # Padding size preserves time step length
groups=pos_embed.conv_pos_groups,
)
self.layer_drop = layer_drop
self.dropout = transformer.attn_layer_dropout # He initialization
std = math.sqrt((4 * (1.0 - self.dropout)) /
(pos_embed.conv_pos * pos_embed.embedding_dim))
nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
nn.init.constant_(self.pos_conv.bias, 0)
self.pos_conv = nn.utils.weight_norm(self.pos_conv,
name="weight",
dim=2)
self.pos_conv = nn.Sequential(self.pos_conv,
SamePad(pos_embed.conv_pos), nn.GELU())
self.layer_norm = nn.LayerNorm(emb_dim)
self.apply(lambda x: transformer_weights_init(x, xavier=False))
def post_init(self, use_transformer_init: bool):
"""
Common post-processing to be called at the end of concrete Classifiers init methods
Args:
use_transformer_init : whether or not to apply transformer_weights_init
"""
if use_transformer_init:
self.apply(
lambda module: transformer_weights_init(module, xavier=False))
def __init__(self, cfg: DictConfig, trainer: Trainer = None):
# shared params for dataset and data loaders
self.dataset_cfg = cfg.dataset
self.tokenizer = get_tokenizer(
tokenizer_name=cfg.language_model.tokenizer,
vocab_file=cfg.language_model.vocab_file,
special_tokens=cfg.language_model.special_tokens,
)
# make vocabulary size divisible by 8 for fast fp16 training
vocab_size = 8 * math.ceil(self.tokenizer.vocab_size / 8)
# init superclass
super().__init__(cfg=cfg, trainer=trainer)
self.embedding_layer = TransformerEmbedding(
vocab_size=vocab_size,
hidden_size=cfg.language_model.hidden_size,
max_sequence_length=cfg.language_model.max_seq_length,
embedding_dropout=cfg.language_model.get("embedding_dropout", 0.0),
learn_positional_encodings=False,
)
self.encoder = TransformerEncoder(
num_layers=cfg.language_model.num_layers,
hidden_size=cfg.language_model.hidden_size,
mask_future=True,
num_attention_heads=cfg.language_model.num_attn_heads,
inner_size=cfg.language_model.inner_size,
ffn_dropout=cfg.language_model.get("ffn_dropout", 0.0),
hidden_act=cfg.language_model.get("inner_activation", "relu"),
attn_score_dropout=cfg.language_model.get("attn_score_dropout", 0.0),
attn_layer_dropout=cfg.language_model.get("attn_layer_dropout", 0.0),
)
self.log_softmax = TokenClassifier(
hidden_size=cfg.language_model.hidden_size, num_classes=vocab_size, log_softmax=True,
)
std_init_range = 1 / math.sqrt(cfg.language_model.hidden_size)
self.apply(lambda module: transformer_weights_init(module, std_init_range))
# tie weights of embedding and softmax matrices
self.log_softmax.mlp.layer0.weight = self.embedding_layer.token_embedding.weight
self.training_loss = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id)
self.validation_loss = SmoothedCrossEntropyLoss(
pad_id=self.tokenizer.pad_id, predict_last_k=self.dataset_cfg.get("predict_last_k", 0),
)
# Optimizer setup needs to happen after all model weights are ready
self.setup_optimization(cfg.optim)
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 __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 = cfg.get("src_language", None)
self.tgt_language = cfg.get("tgt_language", None)
self.multilingual = cfg.get("multilingual", False)
self.multilingual_ids = []
self.encoder_tokenizer_library = cfg.encoder_tokenizer.get(
'library', 'yttm')
self.decoder_tokenizer_library = cfg.decoder_tokenizer.get(
'library', 'yttm')
# 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=self.encoder_tokenizer_library,
encoder_tokenizer_model=cfg.encoder_tokenizer.get(
'tokenizer_model'),
encoder_bpe_dropout=cfg.encoder_tokenizer.get(
'bpe_dropout', 0.0) if cfg.encoder_tokenizer.get(
'bpe_dropout', 0.0) is not None else 0.0,
encoder_model_name=cfg.encoder.get('model_name') if hasattr(
cfg.encoder, 'model_name') else None,
encoder_r2l=cfg.encoder_tokenizer.get('r2l', False),
decoder_tokenizer_library=self.decoder_tokenizer_library,
encoder_tokenizer_vocab_file=cfg.encoder_tokenizer.get(
'vocab_file', None),
decoder_tokenizer_model=cfg.decoder_tokenizer.tokenizer_model,
decoder_bpe_dropout=cfg.decoder_tokenizer.get(
'bpe_dropout', 0.0) if cfg.decoder_tokenizer.get(
'bpe_dropout', 0.0) is not None else 0.0,
decoder_model_name=cfg.decoder.get('model_name') if hasattr(
cfg.decoder, 'model_name') else None,
decoder_r2l=cfg.decoder_tokenizer.get('r2l', False),
)
if self.multilingual:
if isinstance(self.src_language, ListConfig) and isinstance(
self.tgt_language, ListConfig):
raise ValueError(
"cfg.src_language and cfg.tgt_language cannot both be lists. We only support many-to-one or one-to-many multilingual models."
)
elif isinstance(self.src_language, ListConfig):
for lng in self.src_language:
self.multilingual_ids.append(
self.encoder_tokenizer.token_to_id("<" + lng + ">"))
elif isinstance(self.tgt_language, ListConfig):
for lng in self.tgt_language:
self.multilingual_ids.append(
self.encoder_tokenizer.token_to_id("<" + lng + ">"))
else:
raise ValueError(
"Expect either cfg.src_language or cfg.tgt_language to be a list when multilingual=True."
)
if isinstance(self.src_language, ListConfig):
self.tgt_language = [self.tgt_language] * len(
self.src_language)
else:
self.src_language = [self.src_language] * len(
self.tgt_language)
self.source_processor_list = []
self.target_processor_list = []
for src_lng, tgt_lng in zip(self.src_language, self.tgt_language):
src_prcsr, tgt_prscr = self.setup_pre_and_post_processing_utils(
src_lng, tgt_lng)
self.source_processor_list.append(src_prcsr)
self.target_processor_list.append(tgt_prscr)
else:
# After this call, the model will have self.source_processor and self.target_processor objects
self.setup_pre_and_post_processing_utils(self.src_language,
self.tgt_language)
self.multilingual_ids = [None]
# 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)
checkpoint_file = encoder_cfg_dict.pop('checkpoint_file', None)
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', False),
checkpoint_file=checkpoint_file,
)
# 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,
pre_ln_final_layer_norm=decoder_cfg_dict.get(
'pre_ln_final_layer_norm', 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_fn = NLLLoss(ignore_index=self.decoder_tokenizer.pad_id)
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)
def __init__(self, cfg: DictConfig, trainer: Trainer = None):
# 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
# shared params for dataset and data loaders
self.dataset_cfg = cfg.dataset
vocab_file = cfg.language_model.get("vocab_file", None)
if vocab_file is not None:
vocab_file = self.register_artifact("language_model.vocab_file",
vocab_file)
tokenizer_model = cfg.language_model.get("tokenizer_model", None)
if tokenizer_model is not None:
tokenizer_model = self.register_artifact(
"language_model.tokenizer_model", tokenizer_model)
if cfg.language_model.special_tokens:
special_tokens = OmegaConf.to_container(
cfg.language_model.special_tokens, resolve=True)
else:
special_tokens = None
self.tokenizer = get_tokenizer(
tokenizer_name=cfg.language_model.tokenizer,
vocab_file=vocab_file,
special_tokens=special_tokens,
tokenizer_model=tokenizer_model,
)
# make vocabulary size divisible by 8 for fast fp16 training
vocab_size = 8 * math.ceil(self.tokenizer.vocab_size / 8)
# init superclass
super().__init__(cfg=cfg, trainer=trainer)
self.embedding_layer = TransformerEmbedding(
vocab_size=vocab_size,
hidden_size=cfg.language_model.hidden_size,
max_sequence_length=cfg.language_model.max_seq_length,
embedding_dropout=cfg.language_model.get("embedding_dropout", 0.0),
learn_positional_encodings=False,
)
self.encoder = TransformerEncoder(
num_layers=cfg.language_model.num_layers,
hidden_size=cfg.language_model.hidden_size,
mask_future=True,
num_attention_heads=cfg.language_model.num_attn_heads,
inner_size=cfg.language_model.inner_size,
ffn_dropout=cfg.language_model.get("ffn_dropout", 0.0),
hidden_act=cfg.language_model.get("inner_activation", "relu"),
attn_score_dropout=cfg.language_model.get("attn_score_dropout",
0.0),
attn_layer_dropout=cfg.language_model.get("attn_layer_dropout",
0.0),
)
self.log_softmax = TokenClassifier(
hidden_size=cfg.language_model.hidden_size,
num_classes=vocab_size,
log_softmax=True,
)
std_init_range = 1 / math.sqrt(cfg.language_model.hidden_size)
self.apply(
lambda module: transformer_weights_init(module, std_init_range))
# tie weights of embedding and softmax matrices
self.log_softmax.mlp.layer0.weight = self.embedding_layer.token_embedding.weight
if hasattr(self.tokenizer, 'pad_token'):
pad_id = self.tokenizer.pad_id
else:
raise ValueError(
"The tokenizer must support a special `pad_token`. Provide it using"
"the `special_tokens` dictionary.")
self.training_loss = SmoothedCrossEntropyLoss(pad_id=pad_id)
self.validation_loss = SmoothedCrossEntropyLoss(
pad_id=pad_id,
predict_last_k=self.dataset_cfg.get("predict_last_k", 0),
)
self.training_perplexity = Perplexity(dist_sync_on_step=True)
self.validation_perplexity = Perplexity(compute_on_step=False)
# Optimizer setup needs to happen after all model weights are ready
self.setup_optimization()
