hud-python 0.4.43__tar.gz → 0.4.44__tar.gz

{hud_python-0.4.43 → hud_python-0.4.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.4.43
+Version: 0.4.44
 Summary: SDK for the HUD platform.
 Project-URL: Homepage, https://github.com/hud-evals/hud-python
 Project-URL: Bug Tracker, https://github.com/hud-evals/hud-python/issues

{hud_python-0.4.43 → hud_python-0.4.44}/hud/rl/distributed.py

@@ -66,7 +66,13 @@ def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
 
 
 def broadcast_object(obj: Any, src: int = 0) -> Any:
-    """Broadcast a Python object from src rank to all ranks."""
+    """Broadcast a Python object from src rank to all ranks.
+
+    Args:
+        obj: Object to broadcast (used on src rank)
+        src: Source rank
+        device: Device for temporary tensor buffer during pickling transfer
+    """
     if not dist.is_initialized():
         return obj
 
@@ -75,6 +81,33 @@ def broadcast_object(obj: Any, src: int = 0) -> Any:
     return obj_list[0]
 
 
+def scatter_object(
+    obj_list: list[Any] | None,
+    src: int = 0,
+) -> Any:
+    """Scatter a list of Python objects from src so each rank receives one object.
+
+    Usage:
+        - On src rank: pass the full list (length == world_size)
+        - On non-src ranks: pass None
+
+    Returns:
+        The object intended for this rank.
+    """
+    if not dist.is_initialized():
+        # Single-process: return first element if provided, else None
+        if obj_list is None or len(obj_list) == 0:
+            return None
+        return obj_list[0]
+
+    out: list[Any] = [None]
+    if dist.get_rank() == src:
+        dist.scatter_object_list(out, obj_list, src=src)
+    else:
+        dist.scatter_object_list(out, None, src=src)
+    return out[0]
+
+
 def gather_tensors(tensor: torch.Tensor) -> list[torch.Tensor] | None:
     """Gather tensors from all ranks to rank 0.
 

{hud_python-0.4.43 → hud_python-0.4.44}/hud/rl/learner.py

@@ -240,6 +240,8 @@ class GRPOLearner:
                 if sample.inputs:
                     sample = sample.to_device(self.device)
                     sample.old_logprobs, _ = self.compute_logprobs(self.policy, sample.inputs)
+                    # Free GPU memory for this sample immediately
+                    sample.to_device(torch.device("cpu"))
 
             policy_module = self.policy.module if hasattr(self.policy, "module") else self.policy
             with policy_module.disable_adapter():
@@ -247,7 +249,10 @@ class GRPOLearner:
                     if is_main_process():
                         progress.update(f"Processing batch of traces... {i}/{len(batch)}")
                     if sample.inputs:
+                        # Move back to GPU for reference computation, then free
+                        sample = sample.to_device(self.device)
                         sample.ref_logprobs, _ = self.compute_logprobs(self.policy, sample.inputs)
+                        sample.to_device(torch.device("cpu"))
 
         hud_console.info_log("Creating mini-batches...")
         group_size = self.config.training.group_size
@@ -488,15 +493,21 @@ class GRPOLearner:
             out = model(**model_inputs)
 
             logits = out.logits / self.config.actor.temperature
-            log_probs = F.log_softmax(logits, dim=-1)
 
+            # Compute token log-probs via negative cross-entropy to avoid materializing full log_probs
             targets = inputs["input_ids"][:, 1:]
-            token_log_probs = log_probs[:, :-1].gather(-1, targets.unsqueeze(-1)).squeeze(-1)
+            logits_slice = logits[:, :-1, :]
+            loss_flat = F.cross_entropy(
+                logits_slice.reshape(-1, logits_slice.size(-1)),
+                targets.reshape(-1),
+                reduction="none",
+            )
+            token_log_probs = (-loss_flat).reshape_as(targets)
 
             # Compute entropy only for assistant tokens to save memory
             assistant_mask = inputs["assistant_mask"]
             entropy = torch.zeros_like(token_log_probs)
-            if assistant_mask.any():
+            if assistant_mask.any() and getattr(self.config.training, "entropy_beta", 0.0) != 0.0:
                 entropy[assistant_mask] = entropy_from_logits(logits[:, :-1][assistant_mask])
 
             return token_log_probs, entropy
@@ -506,8 +517,20 @@ class GRPOLearner:
             # Return dummy values that match expected shapes
             seq_len = inputs["input_ids"].shape[1] - 1 if "input_ids" in inputs else 0
             batch_size = inputs["input_ids"].shape[0] if "input_ids" in inputs else 1
-            dummy_logprobs = torch.zeros(batch_size, seq_len, device=self.device)
-            dummy_entropy = torch.zeros(batch_size, seq_len, device=self.device)
+            # Create dummy tensors that still participate in autograd so backward doesn't fail
+            try:
+                param_sum = torch.sum(
+                    next(self.policy.parameters())
+                )  # touch params to build a graph
+                base = param_sum * 0.0
+            except StopIteration:
+                base = torch.tensor(0.0, device=self.device)
+            dummy_logprobs = (
+                base + torch.zeros(batch_size, seq_len, device=self.device)
+            ).requires_grad_(True)
+            dummy_entropy = (
+                base + torch.zeros(batch_size, seq_len, device=self.device)
+            ).requires_grad_(True)
             return dummy_logprobs, dummy_entropy
 
     def save(self, path: str) -> None:

{hud_python-0.4.43 → hud_python-0.4.44}/hud/rl/train.py

@@ -13,7 +13,7 @@ import json
 import logging
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import hud
 from hud.rl.actor import Actor
@@ -25,6 +25,7 @@ from hud.rl.distributed import (
     get_global_rank,
     get_world_size,
     is_main_process,
+    scatter_object,
     setup_distributed,
     synchronize,
 )
@@ -133,53 +134,71 @@ async def train(config: Config, tasks: list[Task]) -> None:
             global_reward_stats = None
             global_advantage_stats = None
 
-            # Only rank 0 runs tasks and collects traces
+            # Step-state gate: ensure all ranks branch coherently
+            state = {"ok": False, "err": None, "num_samples": 0}
+            rank_samples = None
+            episode_time_value = None
+
+            # Only rank 0 runs tasks and prepares distribution
             if is_main_process() and actor is not None:
                 import time
 
-                episode_start_time = time.time()
-                traces = await actor.run_tasks(tasks, job_id=job_id)
-                episode_time = time.time() - episode_start_time
-                hud_console.info(f"Sampled {len(traces)} traces in {episode_time:.1f}s")
-                trace_buffer.add(traces)
-                global_reward_stats = [trace.reward for trace in traces]
-
-                # Get all traces from buffer for distribution
-                all_traces = trace_buffer.sample_traces()
-
-                assert len(traces) == len(all_traces)  # noqa: S101
-
-                # Preprocess traces to training samples
-                preprocessed_traces = preprocess_advantages(all_traces, config)
-
-                # Store these for later use in metrics
-                global_advantage_stats = [sample.advantage for sample in preprocessed_traces]
-
-                # Distribute preprocessed samples in groups across ranks
-                gpu_batch_size = len(preprocessed_traces) // num_gpus
-                rank_samples = [
-                    preprocessed_traces[i : i + gpu_batch_size]
-                    for i in range(0, len(preprocessed_traces), gpu_batch_size)
-                ]
+                try:
+                    episode_start_time = time.time()
+                    traces = await actor.run_tasks(tasks, job_id=job_id)
+                    episode_time = time.time() - episode_start_time
+                    hud_console.info(f"Sampled {len(traces)} traces in {episode_time:.1f}s")
+                    trace_buffer.add(traces)
+                    global_reward_stats = [trace.reward for trace in traces]
+
+                    # Get all traces from buffer for distribution
+                    all_traces = trace_buffer.sample_traces()
+
+                    # Preprocess traces to training samples
+                    preprocessed_traces = preprocess_advantages(all_traces, config)
+
+                    # Store these for later use in metrics
+                    global_advantage_stats = [sample.advantage for sample in preprocessed_traces]
+
+                    # Distribute preprocessed samples in groups across ranks via scatter
+                    # Ensure list length is a multiple of num_gpus by allowing empty per-rank slices
+                    gpu_batch_size = max(1, (len(preprocessed_traces) + num_gpus - 1) // num_gpus)
+                    rank_samples = [
+                        preprocessed_traces[i : i + gpu_batch_size]
+                        for i in range(0, len(preprocessed_traces), gpu_batch_size)
+                    ]
+                    # Pad rank_samples to exactly num_gpus entries
+                    if len(rank_samples) < num_gpus:
+                        rank_samples.extend([[] for _ in range(num_gpus - len(rank_samples))])
+
+                    # Log distribution info
+                    dist_msg = (
+                        f"Distributing {len(preprocessed_traces)} samples as {gpu_batch_size} "
+                        f"sized batches across {num_gpus} GPUs"
+                    )
+                    hud_console.info(dist_msg)
+                    for rank in range(num_gpus):
+                        n_samples = len(rank_samples[rank]) if rank < len(rank_samples) else 0
+                        hud_console.info(f"  Rank {rank}: {n_samples} samples")
+
+                    hud_console.section_title(f"Training on {len(all_traces)} traces")
+                    episode_time_value = episode_time
+
+                    state.update({"ok": True, "num_samples": len(preprocessed_traces)})
+                except Exception as e:
+                    state.update({"ok": False, "err": str(e)})
+
+            # Broadcast step-state to keep ranks in lockstep
+            state = broadcast_object(state, src=0)
+            if not state.get("ok", False):
+                hud_console.warning("Step failed on rank 0; skipping this step coherently")
+                synchronize()
+                continue
 
-                # Log distribution info
-                hud_console.info(
-                    f"Distributing {len(preprocessed_traces)} samples as {gpu_batch_size} sized batches across {num_gpus} GPUs"  # noqa: E501
-                )
-                for rank in range(num_gpus):
-                    n_samples = len(rank_samples[rank])
-                    hud_console.info(f"  Rank {rank}: {n_samples} samples")
-
-                hud_console.section_title(f"Training on {len(all_traces)} traces")
-                episode_time_value = episode_time
-            else:
-                rank_samples = None
-                episode_time_value = None
-
-            # Broadcast each rank's samples and episode time
-            rank_samples = broadcast_object(rank_samples, src=0)
+            # Scatter per-rank samples; each rank receives only its slice
+            my_samples = scatter_object(rank_samples if is_main_process() else None, src=0)
+            # Broadcast the episode time (small object)
             episode_time_value = broadcast_object(episode_time_value, src=0)
-            my_samples = rank_samples[get_global_rank()] if rank_samples else []
 
             # Process only assigned samples
             last_metrics = learner.update(my_samples)
@@ -356,7 +375,8 @@ async def main() -> None:
             )
 
     # Run training
-    await train(config, tasks)
+    tasks_typed = cast("list[Task]", tasks)
+    await train(config, tasks_typed)
 
 
 if __name__ == "__main__":

{hud_python-0.4.43 → hud_python-0.4.44}/hud/utils/tests/test_version.py

@@ -5,4 +5,4 @@ def test_import():
     """Test that the package can be imported."""
     import hud
 
-    assert hud.__version__ == "0.4.43"
+    assert hud.__version__ == "0.4.44"

{hud_python-0.4.43 → hud_python-0.4.44}/hud/version.py

@@ -4,4 +4,4 @@ Version information for the HUD SDK.
 
 from __future__ import annotations
 
-__version__ = "0.4.43"
+__version__ = "0.4.44"

{hud_python-0.4.43 → hud_python-0.4.44}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hud-python"
-version = "0.4.43"
+version = "0.4.44"
 description = "SDK for the HUD platform."
 readme = "README.md"
 requires-python = ">=3.11, <3.13"