| --- |
| library_name: transformers |
| base_model: Qwen/Qwen3.5-9B |
| tags: |
| - reinforcement-learning |
| - checkpoint |
| --- |
| |
| # rl-test |
|
|
| Checkpoints from an ongoing, experimental reinforcement-learning run on a Qwen3.5-9B model that writes grounded prompts for short commercial videos. A new iteration is added about every 1.5 to 3.5 hours. Quality has not been evaluated, so treat these as work-in-progress. |
|
|
| ## What is in this repository |
|
|
| | Path | Contents | Use it for | |
| |---|---|---| |
| | `hf/iter_XXXXXXX/` | Hugging Face weights, config, tokenizer and processor (about 18 GB) | **Inference** | |
| | `iter_XXXXXXX/` | Full Megatron `torch_dist` checkpoint with optimizer state (about 136 GB) | Resuming training | |
| | `rollout/global_dataset_state_dict_N.pt` | Data-loader position for iteration N | Resuming training | |
| | `latest_checkpointed_iteration.txt` | Number of the newest full checkpoint | Resuming training | |
|
|
| Iterations are never overwritten. For inference you only need one `hf/iter_XXXXXXX/` folder. Its config, tokenizer, chat template and processor files are identical to the supervised (SFT) checkpoint the run started from. |
|
|
| ## Inference settings |
|
|
| Use the same settings as training. Other settings can change output quality. |
|
|
| | Setting | Value | Note | |
| |---|---|---| |
| | System prompt | Required: `system_prompt.txt`, then the date (`YYYY-MM-DD`) on its own last line | Not in this repository; it is `rl_code/prompts/system_prompt.txt` in the private project dataset. The tool definitions are inside this text. | |
| | `tools=` argument | Do not pass it | | |
| | Thinking | Leave on | Do not pass `enable_thinking=False`. | |
| | Sampling | `do_sample=True`, `temperature=0.7`, `top_p=1.0`, `top_k=0` | `top_k=0` disables top-k; training did not use it. | |
| | `max_new_tokens` | 6144 per turn | | |
| | Total context | 65536 tokens | | |
| | `dtype` | `bfloat16` | | |
| | Images | Downscale only if larger than 4,194,304 pixels (4096 image tokens), keeping aspect ratio | | |
| | Tool loop | At most 8 tool calls and 2 `image_search` calls | Each turn, act on the first `<tool_call>` or `<answer>`. Return tool results as a `user` message wrapped in `<tool_response>`. Put each image right after its `IMG_###: ` label, with at most 5 images per search. | |
|
|
| The stop tokens come from the included `generation_config.json`. The loop itself is implemented as `VideoSearcherEnv` in `rl_code/video_searcher_rl.py` in the private project dataset; use that code if you need the full agent, including its tools and final-message handling. |
|
|
| ## First turn example |
|
|
| Requires `transformers>=5.12`, `torch`, `accelerate` and `huggingface_hub`. Put `system_prompt.txt` in the working directory and set `ITER` to any folder listed under `hf/`. |
|
|
| ```python |
| import time |
| import torch |
| from huggingface_hub import snapshot_download |
| from transformers import AutoModelForImageTextToText, AutoTokenizer |
| |
| ITER = "iter_0000036" |
| BRIEF = "Your video brief here." |
| |
| root = snapshot_download("KL-PHD/rl-test", allow_patterns=[f"hf/{ITER}/*"]) |
| path = f"{root}/hf/{ITER}" |
| tokenizer = AutoTokenizer.from_pretrained(path) |
| model = AutoModelForImageTextToText.from_pretrained(path, dtype=torch.bfloat16, device_map="auto").eval() |
| |
| system = open("system_prompt.txt").read().rstrip("\n") + "\n" + time.strftime("%Y-%m-%d") |
| messages = [{"role": "system", "content": system}, {"role": "user", "content": BRIEF}] |
| inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device) |
| with torch.no_grad(): |
| output = model.generate(**inputs, max_new_tokens=6144, do_sample=True, temperature=0.7, top_p=1.0, top_k=0) |
| print(tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False)) |
| ``` |
|
|
| ## Resuming training |
|
|
| Download one `iter_XXXXXXX/` folder and its `rollout/global_dataset_state_dict_N.pt` into a directory, write `N` to `latest_checkpointed_iteration.txt` there, and pass that directory as `--load`. This needs [slime](https://github.com/THUDM/slime) at commit `540ea91` and 4 GPUs (saved with tensor-parallel size 2). |
|
|
