# deepspeed_config_zero3.yaml
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

python

import os
from trl import RewardTrainer, RewardConfig
from datasets import load_dataset


os.environ["WANDB_PROJECT"] = "xx"

training_args = RewardConfig(
    output_dir="Qwen2.5-7B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter/",
    report_to="wandb",
    run_name="Qwen2.5-7B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter",
    model_init_kwargs={"num_labels": 1}, 
    num_train_epochs=3,
    per_device_train_batch_size=1,  
    gradient_accumulation_steps=64,  
    learning_rate=1e-5,                      
    warmup_ratio=0.1,
    center_rewards_coefficient=1e-2, 
    bf16=True,
    max_length=None,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
)

trainer = RewardTrainer(
    model="../model/Qwen/Qwen2.5-7B-Instruct",
    args=training_args,
    train_dataset=load_dataset("trl-lib/ultrafeedback_binarized", split="train"),
)
trainer.train()

ccelerate launch --config_file deepspeed_config_zero3.yaml train_reward_model_7b.py

Quick start

python

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_PATH = "Qwen2.5-7B-Instruct-ultrafeedback_binarized-reward-num_labels_1"


tokenizer = None
model = None
device = None

def load_model():
    global tokenizer, model, device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Loading model from {MODEL_PATH}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_PATH,
        num_labels=1,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    )
    model.to(device)
    model.eval()
    print(f"Model loaded on {device}")


load_model()
def get_reward_score(prompt, response):
    messages = [
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": response}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    inputs = tokenizer(text, return_tensors="pt", truncation=False, max_length=None).to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
        reward = outputs.logits[0, 0].item()
    
    return reward

Training procedure

Visualize in Weights & Biases

This model was trained with Reward.

Framework versions

TRL: 0.27.0
Transformers: 4.56.2
Pytorch: 2.6.0+cu124
Datasets: 4.4.1
Tokenizers: 0.22.1

Citations

Cite TRL as:

bibtex

@misc{vonwerra2022trl,
	title        = {{TRL: Transformer Reinforcement Learning}},
	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
	year         = 2020,
	journal      = {GitHub repository},
	publisher    = {GitHub},
	howpublished = {\url{https://github.com/huggingface/trl}}
}

chaosc/Qwen2.5-7B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter

Model Card for Qwen2.5-7B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter

Train

Quick start

Training procedure

Framework versions

Citations

chaosc/
Qwen2.5-7B-Instruct-ultrafeedback_binarized-reward-num_labels_1_wo_filter