fix: VLM tensor issue

Author: bradhilton

src/art/serverless/backend.py

@@ -4,7 +4,7 @@
 from openai._types import NOT_GIVEN
 from tqdm import auto as tqdm
 
-from art.client import Client, ExperimentalTrainingConfig
+from art.serverless.client import Client, ExperimentalTrainingConfig
 from art.utils.deploy_model import LoRADeploymentJob, LoRADeploymentProvider
 
 from .. import dev
@@ -128,8 +128,10 @@ async def _train_model(
             model_id=model.id,
             trajectory_groups=trajectory_groups,
             experimental_config=ExperimentalTrainingConfig(
+                advantage_balance=dev_config.get("advantage_balance"),
                 learning_rate=config.learning_rate,
                 precalculate_logprobs=dev_config.get("precalculate_logprobs"),
+                scale_rewards=dev_config.get("scale_rewards"),
             ),
         )
         after: str | None = None

src/art/client.py src/art/serverless/client.pysrc/art/client.py renamed to src/art/serverless/client.py

@@ -19,7 +19,7 @@
 from openai.pagination import AsyncCursorPage
 from typing_extensions import override
 
-from .trajectories import TrajectoryGroup
+from ..trajectories import TrajectoryGroup
 
 ResponseT = TypeVar("ResponseT")
 
@@ -50,8 +50,10 @@ class DeleteCheckpointsResponse(BaseModel):
 
 
 class ExperimentalTrainingConfig(TypedDict, total=False):
+    advantage_balance: float | None
     learning_rate: float | None
     precalculate_logprobs: bool | None
+    scale_rewards: bool | None
 
 
 class TrainingJob(BaseModel):

src/art/unsloth/train.py

@@ -69,11 +69,11 @@ def compute_loss(
                     # if param_group.get("weight_decay"):
                     #     param_group["weight_decay"] = config.weight_decay
 
-        if inputs["pixel_values"][0] is not None:
+        if inputs.get("pixel_values") and inputs["pixel_values"][0] is not None:
             inputs["pixel_values"] = inputs["pixel_values"][0]  # type: ignore
         else:
             del inputs["pixel_values"]  # type: ignore
-        if inputs["image_grid_thw"][0] is not None:
+        if inputs.get("image_grid_thw") and inputs["image_grid_thw"][0] is not None:
             inputs["image_grid_thw"] = inputs["image_grid_thw"][0]  # type: ignore
         else:
             del inputs["image_grid_thw"]  # type: ignore
@@ -114,9 +114,9 @@ def compute_loss(
         next_input_ids = shift_tensor(inputs["tokens"], 0)
         chunk_size = _config.get("logprob_calculation_chunk_size", 1024)
         # Assert that sequence length is evenly divisible by the chunk size
-        assert seq_len % chunk_size == 0, (
-            f"Sequence length ({seq_len}) must be evenly divisible by chunk size ({chunk_size})"
-        )
+        assert (
+            seq_len % chunk_size == 0
+        ), f"Sequence length ({seq_len}) must be evenly divisible by chunk size ({chunk_size})"
         os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
         forward_kwargs = {}
         if "pixel_values" in inputs:
@@ -371,7 +371,9 @@ def _calculate_logprobs(
         chunk_logits = torch.matmul(chunk_hs, lm_head_t)  # [B, chunk_size, V]
         chunk_selected_logits = torch.gather(
             chunk_logits, dim=-1, index=chunk_input_ids.unsqueeze(-1)
-        ).squeeze(-1)  # [B, chunk_size]
+        ).squeeze(
+            -1
+        )  # [B, chunk_size]
         chunk_logsumexp = torch.logsumexp(chunk_logits, dim=-1)  # [B, chunk_size]
         log_probs[:, i : i + chunk_size] = chunk_selected_logits - chunk_logsumexp