def train_step():
    model.eval()
    # 1. completion
    input = ["prompt"] * group_num
    # take vllm as an example
    policy_model_output = vllm_inference(input)  # no grad
    refer_model_output = vllm_inference(input)  # no grad
    # 2. compute reward(advantage)
    rewards = [r(policy_model_output) for r in reward_func_list] # no grad
    # [Optional] reweight reward
    rewards = rewards * reward_reweight

    # 3. compute loss, start traint
    # NOTE: grad required
    model.train()
    kl = compute_kl(policy_model_output, refer_model_output)
    loss = compute_loss(kl, rewards)
    loss.backward()