Issue Reproducing s1.1-32B Training Loss (Observed vs. WandB)

[dzh19990407]

Introduction:
First, thank you for the excellent paper and codebase.

Goal:
I am attempting to reproduce the training results for the s1.1-32B experiment as reported.

My Setup:

• Model: Qwen/Qwen2.5-32B-Instruct
• Key Parameter: block_size=20000
• Hardware: 16 x A100 80G GPUs
• I used train/sft_multinode.sh for training and did not modify any other code.

Problem:
The training loss curve I am observing during my retraining run does not match the curve provided in your primary WandB report for this experiment.

• Expected Behavior (based on WandB): Training loss should decrease steadily and reach a value around 0.4.
• Actual Behavior: My training loss curve follows a trajectory more similar to the example curve shown in the paper's Figure 9. The loss seems to decrease in steps and eventually settles at a lower value (i.e., lower than 0.1).

Request:
Could you clarify if there are known differences in configuration or setup between the run shown in the WandB report (achieving ~0.4 loss) and the conditions that might lead to the curve shape shown in the appendix? Any guidance on replicating the ~0.4 loss result would be appreciated.

Thank you for your time and assistance.

[Muennighoff]

Are you training on the s1K-1.1 dataset with the deepseek traces? Can you maybe share your wandb? The main difference is that they were run on different hardware but that shouldn't make a difference I think.

Also, what evaluation performance are you seeing? It could be that it is just because your data order is different and thus loss looks differently but maybe final perf on aime etc is similar?

[dzh19990407]

Yes, I train with the deepseek traces (i.e., simplescaling/s1K-1.1_tokenized)
I don't know if you can view it. I saved it in my personal wandb profile project.
https://wandb.ai/dingzihan/meta-solver/runs/ujmorsr1?nw=nwuserdingzihan

[dzh19990407]

Evaluation performance is a major concern because the model's output doesn't appear to understand the <|im_start|>answer marker, leading to an infinite generation loop that typically stops only when the maximum token limit is reached. For instance: <|im_start|>answer\nThe final answer is: xxx\n Wait, [continue thinking without end.]

I attach my result files with max_gen_toks=20000,max_tokens_thinking=auto

results.zip

[Muennighoff]

ah yeah i cant see it (you could make it public if you want)

Looks like an issue with the dataset - I would check (e.g. print out) the exact sequences being passed to the model.
For the eval, it looks like maybe it is accidentally using budget forcing in the answer part? Maybe try running it once without the max_tokens_thinking=auto or any other _thinking kwargs i.e. no budget forcing and check whether the answers still look that way?

[dzh19990407]

I will check the test code later based on your suggestion and then proceed to make the workspace publicly accessible. Thanks for your advice!

[dzh19990407]

I have made my workspace public. You can try again at your convenience.

[dzh19990407]

I run the evaluation without thinking kwargs:

lm_eval --model vllm --model_args pretrained=/path/to/my/model,dtype=float32,tensor_parallel_size=8 --tasks aime24_figures --batch_size auto --apply_chat_template --output_path /path/to/output_dir --log_samples --gen_kwargs max_gen_toks=20000

But I found that the model does not seem to output <|im_start|>answer.

results_without_thinking_kwargs.zip

[dzh19990407]

Analysis of the initial model output indicates it's struggling with the <|im_start|>answer portion. Despite budget forcing, the model outputs an answer followed by repetitive reasoning, failing to generate a complete solution. This points to a potential problem with the training's maximum sequence length. Reviewing sft.py, I suspected environment discrepancies (versions of trl, transformers, etc.). I've re-established the environment based on requirements.txt and initiated a new training run. Notably, the loss curve has changed, and training time has increased from ~20-30 minutes to ~2 hours! I believe this addresses the length issue and will share the results after training completion.