revert dtype calls to transformers 4.55.0

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

Author: pstjohn

models/esm2/src/esm/modeling_esm_te.py

@@ -128,7 +128,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.dtype,
+                    params_dtype=config.torch_dtype,
                     window_size=(-1, -1),
                 )
                 for i in range(config.num_hidden_layers)

models/esm2/tests/test_thd_inputs.py

@@ -30,7 +30,9 @@ def te_model_checkpoint(tmp_path):
 
 
 def test_thd_from_collator_output(te_model_checkpoint, input_data_thd):
-    model_thd = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, attn_input_format="thd", dtype=torch.bfloat16)
+    model_thd = NVEsmForMaskedLM.from_pretrained(
+        te_model_checkpoint, attn_input_format="thd", torch_dtype=torch.bfloat16
+    )
     model_thd.to("cuda")
     input_data_thd = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in input_data_thd.items()}
     with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
@@ -76,8 +78,10 @@ def test_thd_values_match(te_model_checkpoint, tokenizer):
     input_data_bhsd = bhsd_collator(sequences)
     input_data_thd = thd_collator(sequences)
 
-    model_bshd = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, dtype=torch.bfloat16)
-    model_thd = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, attn_input_format="thd", dtype=torch.bfloat16)
+    model_bshd = NVEsmForMaskedLM.from_pretrained(te_model_checkpoint, torch_dtype=torch.bfloat16)
+    model_thd = NVEsmForMaskedLM.from_pretrained(
+        te_model_checkpoint, attn_input_format="thd", torch_dtype=torch.bfloat16
+    )
     model_bshd.to("cuda")
     model_thd.to("cuda")