Switch deprecated torch_dtype parameter to dtype

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

models/README.md

@@ -111,7 +111,7 @@ def convert_hf_to_te(model_hf: nn.Module, **config_kwargs) -> nn.Module:
     """Convert HuggingFace model to TransformerEngine format."""
     te_config = MyModelTEConfig(**model_hf.config.to_dict(), **config_kwargs)
     with init_empty_weights():
-        model_te = MyModelTE(te_config, torch_dtype=te_config.torch_dtype)
+        model_te = MyModelTE(te_config, dtype=te_config.dtype)
 
     output_model = io.apply_transforms(model_hf, model_te, ...)
     return output_model

models/amplify/export.py

@@ -36,7 +36,7 @@
     # Smoke test that the model can be loaded.
     model_te = AutoModelForMaskedLM.from_pretrained(
         f"./checkpoint_export/{tag}",
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,
         trust_remote_code=True,
     )
     del model_te

models/amplify/src/amplify/amplify_te.py

@@ -147,17 +147,15 @@ def __init__(self, config: AMPLIFYConfig, **kwargs):
             config.padded_vocab_size,
             config.hidden_size,
             padding_idx=config.pad_token_id,
-            dtype=config.torch_dtype,
+            dtype=config.dtype,
         )
 
         if config.layer_norm_after_embedding:
             self.layer_norm_1 = (
-                transformer_engine.pytorch.RMSNorm(
-                    config.hidden_size, config.norm_eps, params_dtype=config.torch_dtype
-                )
+                transformer_engine.pytorch.RMSNorm(config.hidden_size, config.norm_eps, params_dtype=config.dtype)
                 if config.rms_norm
                 else transformer_engine.pytorch.LayerNorm(
-                    config.hidden_size, config.norm_eps, params_dtype=config.torch_dtype
+                    config.hidden_size, config.norm_eps, params_dtype=config.dtype
                 )
             )
 
@@ -194,7 +192,7 @@ def __init__(self, config: AMPLIFYConfig, **kwargs):
                     window_size=(-1, -1),
                     rotary_pos_interleaved=True,
                     seq_length=config.max_length,
-                    params_dtype=config.torch_dtype,
+                    params_dtype=config.dtype,
                 )
             )
 
@@ -277,7 +275,7 @@ def __init__(self, config: AMPLIFYConfig, **kwargs):
                 config.hidden_size,
                 config.padded_vocab_size,
                 config.norm_eps,
-                params_dtype=config.torch_dtype,
+                params_dtype=config.dtype,
                 normalization="RMSNorm" if config.rms_norm else "LayerNorm",
                 init_method=lambda x: torch.nn.init.uniform_(
                     x, -self.config.decoder_init_range, self.config.decoder_init_range
@@ -286,7 +284,7 @@ def __init__(self, config: AMPLIFYConfig, **kwargs):
 
         else:
             self.decoder = transformer_engine.pytorch.Linear(
-                config.hidden_size, config.vocab_size, params_dtype=config.torch_dtype
+                config.hidden_size, config.vocab_size, params_dtype=config.dtype
             )
 
     def forward(

models/amplify/src/amplify/state_dict_convert.py

@@ -46,7 +46,7 @@ def convert_amplify_hf_to_te(model_hf: nn.Module, **config_kwargs) -> nn.Module:
     """
     te_config = AMPLIFYConfig(**model_hf.config.to_dict(), **config_kwargs)
     with init_empty_weights():
-        model_te = AMPLIFYForMaskedLM(te_config, torch_dtype=te_config.torch_dtype)
+        model_te = AMPLIFYForMaskedLM(te_config, dtype=te_config.dtype)
 
     output_model = io.apply_transforms(
         model_hf,

models/amplify/tests/conftest.py

@@ -36,7 +36,7 @@ def tokenizer():
 @pytest.fixture
 def config():
     config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True)
-    config.torch_dtype = torch.bfloat16
+    config.dtype = torch.bfloat16
     return config
 
 

models/amplify/tests/test_encoder_block.py

@@ -57,7 +57,7 @@ def data(self) -> torch.Tensor:
 @pytest.fixture
 def config():
     config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True)
-    config.torch_dtype = torch.bfloat16
+    config.dtype = torch.bfloat16
     return config
 
 
@@ -169,7 +169,7 @@ def test_encoder_block_forward(inputs, config):
         window_size=(-1, -1),
         rotary_pos_interleaved=True,
         seq_length=config.max_length,
-        params_dtype=config.torch_dtype,
+        params_dtype=config.dtype,
     ).to("cuda", dtype=torch.bfloat16)
 
     state_dict_mapping = {

models/esm2/src/esm/export.py

@@ -64,7 +64,7 @@ def export_hf_checkpoint(tag: str, export_path: Path):
     # Smoke test that the model can be loaded.
     model_te = AutoModelForMaskedLM.from_pretrained(
         export_path / tag,
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,
         trust_remote_code=True,
     )
     del model_te

models/esm2/src/esm/modeling_esm_te.py

@@ -129,7 +129,7 @@ def __init__(self, config: NVEsmConfig):
                     micro_batch_size=config.micro_batch_size,
                     num_gqa_groups=config.num_attention_heads,
                     fuse_qkv_params=config.fuse_qkv_params,
-                    params_dtype=config.torch_dtype,
+                    params_dtype=config.dtype,
                     window_size=(-1, -1),
                 )
                 for i in range(config.num_hidden_layers)

models/esm2/tests/test_distributed_strategies.py

@@ -145,14 +145,14 @@ def run_forward_backward(use_te: bool, strategy: Strategy, input_data: dict, dis
         if use_te:
             model = AutoModelForMaskedLM.from_pretrained(
                 "nvidia/esm2_t6_8M_UR50D",
-                torch_dtype=torch.bfloat16,
+                dtype=torch.bfloat16,
                 trust_remote_code=True,
             )
             transformer_layers = model.esm.encoder.layers
         else:
             model = AutoModelForMaskedLM.from_pretrained(
                 "facebook/esm2_t6_8M_UR50D",
-                torch_dtype=torch.bfloat16,
+                dtype=torch.bfloat16,
             )
             transformer_layers = model.esm.encoder.layer
             del model.esm.contact_head  # Unused in backwards pass.

recipes/amplify_accelerate_te_fp8/train.py

@@ -41,7 +41,7 @@ def main(args: DictConfig):
     model = AutoModelForMaskedLM.from_config(
         config,
         trust_remote_code=True,
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,
     )
 
     train_dataset, eval_dataset, data_collator = create_datasets_and_collator(