Add UploadCheckpointsCallback (transformers v5 removed PushToHubCallback; periodic saves were local-only)

train_lora.py

@@ -18,11 +18,13 @@ from the harness's own ATIF trace.
 """
 
 import os
+import threading
 from functools import partial
 
 from datasets import Dataset
+from huggingface_hub import HfApi
 from peft import LoraConfig
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, TrainerCallback
 from trl.experimental.async_grpo.async_grpo_config import AsyncGRPOConfig
 from trl.experimental.async_grpo.async_grpo_trainer import AsyncGRPOTrainer
 from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker
@@ -41,6 +43,47 @@ TOOL_CALL_BUDGET = 32   # reference for the efficiency bonus; tune to the task p
 MAX_BONUS = 0.1
 
 
+class UploadCheckpointsCallback(TrainerCallback):
+    """Push every checkpoint to the Hub, in the background.
+
+    transformers v5 removed PushToHubCallback: periodic _save_checkpoint writes
+    only to local disk (save_model is called with _internal_call=True, which
+    skips the push), and `hub_strategy="every_save"` is vestigial. Without this
+    callback a crashed job loses the whole run.
+    """
+
+    def __init__(self, repo_id):
+        self.repo_id = repo_id
+        self._threads = []
+
+    def on_save(self, args, state, control, **kwargs):
+        if not (self.repo_id and state.is_world_process_zero):
+            return control
+        ckpt = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
+        t = threading.Thread(target=self._upload, args=(ckpt,), daemon=True)
+        t.start()
+        self._threads.append(t)
+        return control
+
+    def _upload(self, ckpt):
+        try:
+            HfApi().upload_folder(
+                repo_id=self.repo_id,
+                folder_path=ckpt,
+                path_in_repo=os.path.basename(ckpt),
+                # Keep adapter + trainer state, skip multi-GB optimizer shards.
+                ignore_patterns=["optimizer.pt", "optimizer/**", "*.pth"],
+            )
+            print(f"[hub] uploaded {ckpt}", flush=True)
+        except Exception as exc:  # never let an upload hiccup kill training
+            print(f"[hub] checkpoint upload failed: {exc}", flush=True)
+
+    def on_train_end(self, args, state, control, **kwargs):
+        for t in self._threads:
+            t.join(timeout=600)
+        return control
+
+
 def correctness_plus_efficiency(outcome):
     """1/0 correctness, plus <= 0.1 for a correct answer in fewer tool calls.
 
@@ -124,12 +167,9 @@ def main():
         "seed": 42,
     }
     if ADAPTER_REPO:
-        # Jobs are ephemeral: push every checkpoint (adapter-only) to the Hub.
-        config_kwargs.update(
-            push_to_hub=True,
-            hub_model_id=ADAPTER_REPO,
-            hub_strategy="every_save",
-        )
+        # push_to_hub covers the END-of-training push; periodic checkpoints are
+        # pushed by UploadCheckpointsCallback (transformers v5 dropped that path).
+        config_kwargs.update(push_to_hub=True, hub_model_id=ADAPTER_REPO)
     config = AsyncGRPOConfig(**config_kwargs)
 
     # LoRA on attention + MLP projections. Module names verified against the
@@ -150,6 +190,7 @@ def main():
         processing_class=tokenizer,
         rollout_worker=worker,
         peft_config=peft_config,
+        callbacks=[UploadCheckpointsCallback(ADAPTER_REPO)],
     )
     trainer.train()
     trainer.save_model("./checkpoints/final")   # adapter_model.safetensors + adapter_config.json