The tiny_nemotron5_h training setups in tests/unit/models/policy/test_dtensor_worker.py::TestTwoGPUCluster::test_dtensor_worker_training fail under DTensor v2 when loading the HF checkpoint:
RuntimeError: Missing key in checkpoint state_dict: backbone.embeddings.weight
The model built by AutoModel names the embedding backbone.embeddings.weight (plural), while the checkpoint produced by transformers' NemotronHForCausalLM uses backbone.embedding.weight (singular). All 25 other keys match.
nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py:132 only maps backbone.embed_tokens.weight -> backbone.embeddings.weight; nothing handles the singular spelling.
Note that the underlying error is masked in CI by TypeError: cannot pickle code objects -- _DistWrapper.reduce_scatter gathers the exception object, whose traceback frames hold code objects.
These two setups are disabled in the test until this is fixed.
The
tiny_nemotron5_htraining setups intests/unit/models/policy/test_dtensor_worker.py::TestTwoGPUCluster::test_dtensor_worker_trainingfail under DTensor v2 when loading the HF checkpoint:The model built by AutoModel names the embedding
backbone.embeddings.weight(plural), while the checkpoint produced by transformers'NemotronHForCausalLMusesbackbone.embedding.weight(singular). All 25 other keys match.nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py:132only mapsbackbone.embed_tokens.weight->backbone.embeddings.weight; nothing handles the singular spelling.Note that the underlying error is masked in CI by
TypeError: cannot pickle code objects--_DistWrapper.reduce_scattergathers the exception object, whose traceback frames hold code objects.These two setups are disabled in the test until this is fixed.