fix: fix the deprecated usage of tuple slicing, wandb's start_method, and transformers's dtype (#462)

Author: rchardx

areal/engine/base_hf_engine.py

@@ -154,7 +154,7 @@ def create_device_model(self):
                 model = AutoModelForImageTextToText.from_pretrained(
                     pretrained_model_name_or_path=self.config.path,
                     trust_remote_code=True,
-                    torch_dtype=dtype,
+                    dtype=dtype,
                     attn_implementation=self.config.attn_impl,
                 )
                 if self.config.disable_dropout:
@@ -179,40 +179,35 @@ def create_device_model(self):
 
     def _create_llm_actor_or_critic(self):
         dtype = getattr(torch, self.config.dtype)
-        if not self.config.is_critic:
-            if self.config.init_from_scratch:
-                # initialize model from config
-                # NOTE: VLM cannot directly load state dict using this
-                # random initialized model, so otherwise we call
-                # from_pretrained rather than loading weights into this random model.
-                model = AutoModelForCausalLM.from_config(
-                    self.model_config,
-                    torch_dtype=dtype,
-                    attn_implementation=self.config.attn_impl,
-                )
-            else:
-                model = AutoModelForCausalLM.from_pretrained(
-                    pretrained_model_name_or_path=self.config.path,
-                    trust_remote_code=True,
-                    torch_dtype=dtype,
-                    attn_implementation=self.config.attn_impl,
-                )
+
+        if self.config.is_critic:
+            model_class = AutoModelForTokenClassification
+            model_kwargs = {"num_labels": 1}
         else:
-            if self.config.init_from_scratch:
-                model = AutoModelForTokenClassification.from_config(
-                    self.model_config,
-                    torch_dtype=dtype,
-                    num_labels=1,
-                    attn_implementation=self.config.attn_impl,
-                )
-            else:
-                model = AutoModelForTokenClassification.from_pretrained(
-                    pretrained_model_name_or_path=self.config.path,
-                    trust_remote_code=True,
-                    torch_dtype=dtype,
-                    num_labels=1,
-                    attn_implementation=self.config.attn_impl,
-                )
+            model_class = AutoModelForCausalLM
+            model_kwargs = {}
+
+        common_kwargs = {
+            "dtype": dtype,
+            "attn_implementation": self.config.attn_impl,
+        }
+        model_kwargs.update(common_kwargs)
+
+        if self.config.init_from_scratch:
+            # initialize model from config
+            # NOTE: VLM cannot directly load state dict using this
+            # random initialized model, so otherwise we call
+            # from_pretrained rather than loading weights into this random model.
+            model = model_class.from_config(
+                self.model_config,
+                **model_kwargs,
+            )
+        else:
+            model = model_class.from_pretrained(
+                pretrained_model_name_or_path=self.config.path,
+                trust_remote_code=True,
+                **model_kwargs,
+            )
         return model
 
     def destroy(self):

areal/utils/stats_logger.py

@@ -65,7 +65,6 @@ def init(self):
             force=True,
             id=f"{self.config.experiment_name}_{self.config.trial_name}_{suffix}",
             resume="allow",
-            settings=wandb.Settings(start_method="fork"),
         )
 
         swanlab_config = self.config.swanlab

areal/utils/ulysses.py

@@ -100,7 +100,7 @@ def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor:
         return x
     slc = [slice(None)] * len(x.shape)
     slc[dim] = slice(0, -padding_size)
-    return x[slc]
+    return x[tuple(slc)]
 
 
 def slice_input_tensor(
@@ -118,7 +118,7 @@ def slice_input_tensor(
     parts = x.size(dim) // sp_world_size
     slc = [slice(None)] * len(x.shape)
     slc[dim] = slice(sp_rank * parts, (sp_rank + 1) * parts)
-    return x[slc].contiguous()
+    return x[tuple(slc)].contiguous()
 
 
 def all_to_all_tensor(

evaluation/model_utils.py

@@ -228,7 +228,7 @@ def load_hf_lm_and_tokenizer(
         # defaul load in float16
         model = AutoModelForCausalLM.from_pretrained(
             model_name_or_path,
-            torch_dtype=torch.float16,
+            dtype=torch.float16,
             device_map=device_map,
             trust_remote_code=True,
             use_safetensors=use_safetensors,

examples/docs/debug/cmp_rollout.py

@@ -27,7 +27,7 @@ def image2base64(images: List[ImageObject] | ImageObject) -> List[str] | str:
 model_id = "google/gemma-3-4b-it"
 
 model = Gemma3ForConditionalGeneration.from_pretrained(
-    model_id, torch_dtype=torch.bfloat16
+    pretrained_model_name_or_path=model_id, dtype=torch.bfloat16
 ).to("cuda")
 model.eval()