Spaces:
Sleeping
Sleeping
File size: 1,630 Bytes
aeba3f1 e3d939d aeba3f1 e3d939d aeba3f1 e3d939d aeba3f1 e3d939d aeba3f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Minimal GRPO training entrypoint scaffold.
This file intentionally does not start training on import. It validates that the
required TRL/Trackio configuration can be constructed when optional training
dependencies are installed.
"""
from __future__ import annotations
import os
from training.trackio_utils import build_run_name, get_git_sha
def build_grpo_config():
from trl import GRPOConfig
model_name = os.getenv("MODEL_NAME", "Qwen/Qwen3-1.7B")
difficulty = int(os.getenv("DIFFICULTY", "0"))
output_dir = os.getenv("OUTPUT_DIR", "CyberSecurity_OWASP-qwen3-1.7b-grpo")
trackio_space_id = os.getenv("TRACKIO_SPACE_ID", output_dir)
os.environ.setdefault("TRACKIO_PROJECT", "CyberSecurity_OWASP-grpo")
run_name = os.getenv(
"RUN_NAME",
build_run_name(model_name, "grpo", difficulty, git_sha=get_git_sha()),
)
return GRPOConfig(
output_dir=output_dir,
report_to="trackio",
trackio_space_id=trackio_space_id,
run_name=run_name,
logging_steps=1,
save_steps=25,
learning_rate=5e-6,
num_train_epochs=1,
per_device_train_batch_size=1,
gradient_accumulation_steps=32,
num_generations=2,
max_prompt_length=4096,
max_completion_length=768,
use_vllm=True,
vllm_mode="colocate",
vllm_gpu_memory_utilization=0.2,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={"use_reentrant": False},
push_to_hub=False,
)
def main():
config = build_grpo_config()
print(config)
if __name__ == "__main__":
main()
|