fix: tied embedding v4 to v5

[akoumpa]

What does this PR do ?

class NemotronFlashForCausalLM(NemotronFlashPreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]  # <- here

    def __init__(self, config: NemotronFlashConfig):
        super().__init__(config)
        self.config = config
        self.model = NemotronFlashModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.post_init()

In transformers v4, _tied_weights_keys was used to denote which layers are tied the the nn.Embedding.

In v5 the concept was advanced and instead of being embedding to all, it is now a dictionary with a source and destination FQNs.

As a result, when trying to restore v4 checkpoints in v5, these will fail, since the expectation is that the _tied_weights_keys will be a dictionary.

This PR fixes on-the-fly those cases so that they can be loaded correctly in v5 (previously crashing).

distributed:
  strategy: fsdp2
  dp_size: 8
  tp_size: 1
  cp_size: 1
  pp_size: 1
  sequence_parallel: false
  activation_checkpointing: true
checkpoint:
  checkpoint_dir: checkpoints/
  enabled: false
  model_save_format: torch_save
  save_consolidated: false
dataloader:
  _target_: torchdata.stateful_dataloader.stateful_dataloader.StatefulDataLoader
  collate_fn: nemo_automodel.components.datasets.utils.default_collater
  shuffle: false
dataset:
  _target_: nemo_automodel.components.datasets.llm.squad.make_squad_dataset
  dataset_name: rajpurkar/squad
  split: train
dist_env:
  backend: nccl
  timeout_minutes: 1
loss_fn:
  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: /akoumparouli/automodel/ckpt/nvidia/Nemotron-Flash-1B/
  torch_dtype: bf16
  trust_remote_code: true
  dtype: bfloat16
optimizer:
  _target_: torch.optim.adamw.AdamW
  lr: 0.0002
  weight_decay: 0.1
packed_sequence:
  packed_sequence_size: 512
  split_across_pack: false
step_scheduler:
  ckpt_every_steps: 2000
  global_batch_size: 128
  local_batch_size: 4
  max_steps: 20
  num_epochs: 1

Note: this model has another issue, when initializing on meta the A_log is initialized as float32 while the rest bfloat16, however, the public checkpoint uses bfloat16 everywhere. The issue is fsdp2 cannot wrap a module with mixed dtypes, however, the issue is that the dtype of the checkpoint is not preserved.

model.layers.17.ffn.gate_proj.weight torch.bfloat16
model.layers.17.ffn.down_proj.weight torch.bfloat16
model.layers.17.ffn.up_proj.weight torch.bfloat16
model.layers.17.pre_ffn_layernorm.weight torch.bfloat16
model.layers.18.mamba.dt_bias torch.bfloat16
model.layers.18.mamba.A_log torch.float32
model.layers.18.mamba.D torch.bfloat16
model.layers.18.mamba.in_proj.weight torch.bfloat16
model.layers.18.mamba.conv1d.weight torch.bfloat16
model.layers.18.mamba.conv1d.bias torch.bfloat16
model.layers.18.mamba.norm.weight torch.bfloat16
model.layers.18.mamba.out_proj.weight torch.bfloat16
model.layers.18.input_layernorm.weight torch.bfloat16
model.layers.19.ffn.gate_proj.weight torch.bfloat16

Changelog

• Add specific line by line info of high level changes in this PR.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test a38292d

[akoumpa]

/claude review

[claude]

LGTM

[akoumpa]

/ok to test 1fe2c82

[akoumpa]

/claude review