Install cut-cross-entropy and log CUDA memory

__pycache__/zerogpu_tinker_service.cpython-312.pyc

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:39882bb35d1044aee21117e4e84b63b50fe8a7c12e427d42f3a897a4e36701d7
-size 102030
+oid sha256:786870ce3b5415a2fa92c73c0e1e54ffc685353d2e06afbad45b5c778778ab0d
+size 121790

requirements.txt

@@ -1,5 +1,6 @@
 transformers @ git+https://github.com/huggingface/transformers.git
 hatchery-core[zerogpu] @ git+https://github.com/axolotl-ai-cloud/hatchery-core.git@914c893
+cut-cross-entropy @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git
 torchao
 dflash @ git+https://github.com/winglian/dflash.git@904b558
 flash-linear-attention

train_grpo_local.py

@@ -344,6 +344,13 @@ def log_rollout(wandb_run: Any | None, record: dict[str, Any]) -> None:
                 "rollout/generated_tokens_server": profile.get("generated_tokens"),
             }
         )
+        for key in (
+            "cuda_memory_allocated_bytes",
+            "cuda_memory_reserved_bytes",
+            "cuda_max_memory_allocated_bytes",
+            "cuda_max_memory_reserved_bytes",
+        ):
+            payload[f"rollout/{key}"] = profile.get(key)
     wandb_run.log(payload)
 
 
@@ -390,6 +397,13 @@ def log_train_step(wandb_run: Any | None, record: Any) -> None:
                 "train/cuda_synchronize_ms": phase_duration_ms(profile, "cuda_synchronize"),
             }
         )
+        for key in (
+            "cuda_memory_allocated_bytes",
+            "cuda_memory_reserved_bytes",
+            "cuda_max_memory_allocated_bytes",
+            "cuda_max_memory_reserved_bytes",
+        ):
+            payload[f"train/{key}"] = profile.get(key)
         collate_ms = 0.0
         collate_found = False
         fwd_bwd_ms = 0.0

train_sft_local.py

@@ -148,6 +148,13 @@ def log_train_step(wandb_run: Any | None, record: Any) -> None:
                 "train/cuda_synchronize_ms": phase_duration_ms(profile, "cuda_synchronize"),
             }
         )
+        for key in (
+            "cuda_memory_allocated_bytes",
+            "cuda_memory_reserved_bytes",
+            "cuda_max_memory_allocated_bytes",
+            "cuda_max_memory_reserved_bytes",
+        ):
+            payload[f"train/{key}"] = profile.get(key)
         collate_ms = 0.0
         collate_found = False
         fwd_bwd_ms = 0.0

zerogpu_tinker_service.py

@@ -526,6 +526,11 @@ class _Profiler:
         self.context = {key: value for key, value in context.items() if value is not None}
         self.started = time.monotonic()
         self.phases: list[dict[str, Any]] = []
+        if torch.cuda.is_available():
+            try:
+                torch.cuda.reset_peak_memory_stats()
+            except Exception:
+                pass
 
     @contextmanager
     def phase(self, name: str, **fields: Any):
@@ -545,13 +550,28 @@ class _Profiler:
         _log_event("zerogpu_profile", **self.snapshot(**fields))
 
     def snapshot(self, **fields: Any) -> dict[str, Any]:
-        return {
+        payload = {
             "operation": self.operation,
             "duration_ms": round((time.monotonic() - self.started) * 1000, 2),
             "phases": self.phases,
             **self.context,
             **{key: value for key, value in fields.items() if value is not None},
         }
+        if torch.cuda.is_available():
+            try:
+                payload.update(
+                    {
+                        "cuda_memory_allocated_bytes": int(torch.cuda.memory_allocated()),
+                        "cuda_memory_reserved_bytes": int(torch.cuda.memory_reserved()),
+                        "cuda_max_memory_allocated_bytes": int(
+                            torch.cuda.max_memory_allocated()
+                        ),
+                        "cuda_max_memory_reserved_bytes": int(torch.cuda.max_memory_reserved()),
+                    }
+                )
+            except Exception:
+                pass
+        return payload
 
 
 def _response_context(result: Any) -> dict[str, Any]:
@@ -569,6 +589,18 @@ def _response_context(result: Any) -> dict[str, Any]:
     return context
 
 
+def _fused_ce_kernels_available() -> list[str]:
+    from hatchery.core.fused_losses import _try_import_cce, _try_import_liger
+
+    kernels: list[str] = []
+    if _try_import_cce() is not None:
+        kernels.append("cce")
+    if _try_import_liger() is not None:
+        kernels.append("liger")
+    kernels.append("chunked")
+    return kernels
+
+
 def _write_session_metadata(session: ServiceSession) -> None:
     path = _metadata_path(session.model_id)
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -1866,6 +1898,7 @@ def build_tinker_api() -> FastAPI:
             "max_rank": _supported_rank(),
             "futures_api": True,
             "sdk_compat_version": "tinker-0.x",
+            "fused_ce_kernels_available": _fused_ce_kernels_available(),
             "zero_gpu": True,
             "speculative_decoding": {
                 "backend": "dflash",
@@ -2036,6 +2069,7 @@ def build_space_app() -> Any:
             "max_rank": _supported_rank(),
             "futures_api": True,
             "sdk_compat_version": "tinker-0.x",
+            "fused_ce_kernels_available": _fused_ce_kernels_available(),
             "zero_gpu": True,
             "speculative_decoding": {
                 "backend": "dflash",