Upload folder using huggingface_hub

train_demo.py

@@ -13,6 +13,7 @@ import hashlib
 import re
 import os
 
+import wandb
 from datasets import Dataset
 from trl import GRPOConfig, GRPOTrainer
 from trl.experimental.openenv import generate_rollout_completions
@@ -23,7 +24,7 @@ from skill_invocation_env.client import SkillInvocationEnv
 from skill_invocation_env.models import SkillInvocationAction
 
 # ── Configuration ──────────────────────────────────────────────────────────────
-MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
+MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")
 ENV_URL = os.getenv("ENV_URL", "https://mpnikhil-skill-invocation-env.hf.space")
 HF_TOKEN = os.getenv("HF_TOKEN")
 OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs/qwen-skill-env")
@@ -264,12 +265,18 @@ def rollout_func(prompts: list[str], trainer: GRPOTrainer) -> dict[str, list]:
         seed = _extract_seed(prompt_text)
 
         env = SkillInvocationEnv(base_url=ENV_URL, connect_timeout_s=60)
-        episode = rollout_once(
-            trainer=trainer,
-            env=env,
-            tokenizer=tokenizer,
-            env_seed=seed,
-        )
+        try:
+            episode = rollout_once(
+                trainer=trainer,
+                env=env,
+                tokenizer=tokenizer,
+                env_seed=seed,
+            )
+        finally:
+            try:
+                env.close()
+            except Exception:
+                pass
         all_prompt_ids.append(episode["prompt_ids"])
         all_completion_ids.append(episode["completion_ids"])
         all_logprobs.append(episode["logprobs"])
@@ -363,12 +370,12 @@ if __name__ == "__main__":
         output_dir=OUTPUT_DIR,
         use_vllm=True,
         vllm_mode="colocate",
-        vllm_gpu_memory_utilization=0.6,
+        vllm_gpu_memory_utilization=0.3,
         num_train_epochs=1,
         num_generations=NUM_GENERATIONS,
-        max_completion_length=MAX_COMPLETION_LENGTH,
-        per_device_train_batch_size=8,
-        gradient_accumulation_steps=4,
+        max_completion_length=512,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=32,
         learning_rate=1e-6,
         logging_steps=1,
         save_steps=50,