Search-R1 学习指北

【系统提示词】
你是一个严谨的助手。请根据以下给出的【参考资料】来回答用户的【问题】。
如果参考资料中没有相关信息，请直接回答“不知道”，不要胡思乱想。

【参考资料开始】
资料 [1]: 蒂姆·库克（Tim Cook），1960年出生，现任苹果公司CEO...
资料 [2]: 奥本大学（Auburn University）成立于1856年，位于美国阿拉巴马州...
资料 [3]: 苹果公司由史蒂夫·乔布斯等人于1976年创立...
【参考资料结束】

【用户问题】
苹果现任CEO的母校是哪年建校的？

<think>
用户想知道苹果现任 CEO 母校的建校时间。
第一步：我需要先确定苹果现任 CEO 是谁。
-> 发起搜索：【苹果现任 CEO】
-> 收到结果：蒂姆·库克（Tim Cook）。

第二步：我知道名字了，接下来要查他的母校。
-> 发起搜索：【蒂姆·库克 母校 大学】
-> 收到结果：他毕业于奥本大学（Auburn University）。

第三步：最后查这所大学的建校时间。
-> 发起搜索：【奥本大学 建校时间】
-> 收到结果：1856年。

第四步：信息完整，可以回答。
</think>
苹果现任 CEO 蒂姆·库克的母校是奥本大学，该校建于 1856 年。

data = {
    "data_source": data_source,
    "prompt": [{
        "role": "user",
        "content": question,
    }],
    "ability": "fact-reasoning",
    "reward_model": {
        "style": "rule",
        "ground_truth": {
            "target": example["golden_answers"]
        }
    },
    "extra_info": {
        "split": split,
        "index": idx,
    }
}

def make_prefix(dp, template_type):
    question = dp['question']

    if template_type == 'base':
        """This works for any base model"""
        prefix = f"""Answer the given question. \
You must conduct reasoning inside <think> and </think> first every time you get new information. \
After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search> and it will return the top searched results between <information> and </information>. \
You can search as many times as your want. \
If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. For example, <answer> Beijing </answer>. Question: {question}\n"""
    else:
        raise NotImplementedError
    return prefix

def make_map_fn(split):

        def process_fn(example, idx):
            example['question'] = example['question'].strip()
            if example['question'][-1] != '?':
                example['question'] += '?'
            question = make_prefix(example, template_type=args.template_type)
            # ...

data.train_files=$DATA_DIR/train.parquet
data.val_files=$DATA_DIR/test.parquet

final_gen_batch_output = generation_manager.run_llm_loop(
    gen_batch=gen_batch,
    initial_input_ids=first_input_ids,
)

if data_source in [
    'nq', 'triviaqa', 'popqa', 'web_questions',
    'hotpotqa', '2wikimultihopqa', 'musique',
    'bamboogle', 'strategyqa'
]:
    return qa_em_format.compute_score_em

@app.post("/retrieve")
def retrieve_endpoint(request: QueryRequest):
    results, scores = retriever.batch_search(
        query_list=request.queries,
        num=request.topk,
        return_score=request.return_scores
    )
    # 格式化返回
    ...

{
  "queries": ["What is the capital of France?", "Who wrote Dune?"],
  "topk": 3,
  "return_scores": true
}

{
  "result": [
    [
      {"document": {"title": "France", "text": "...", "contents": "\"France\"\n..."}, "score": 0.95},
      {"document": {"title": "Paris", "text": "...", "contents": "\"Paris\"\n..."}, "score": 0.88},
      ...
    ],
    [...]
  ]
}

class BM25Retriever(BaseRetriever):
    def __init__(self, config):
        from pyserini.search.lucene import LuceneSearcher
        self.searcher = LuceneSearcher(self.index_path)

    def _search(self, query, num, return_score):
        hits = self.searcher.search(query, num)
        # 从 hits 中提取文档内容
        all_contents = [
            json.loads(self.searcher.doc(hit.docid).raw())['contents']
            for hit in hits
        ]
        results = [
            {'title': content.split("\n")[0], 'text': "\n".join(content.split("\n")[1:]), 'contents': content}
            for content in all_contents
        ]
        return results, scores

┌────────────────────────────────────────────────────────────────────────────────────┐
│                           Multi-Turn Rollout Loop                                  │
├────────────────────────────────────────────────────────────────────────────────────┤
│                                                                                    │
│  ┌──────────────┐     ┌──────────────┐     ┌──────────────┐                        │
│  │    Prompt    │────▶│   LLM 生成    │────▶│  后处理响应   │                        │
│  │  (rolling)   │     │ generate_seq │     │ postprocess  │                        │
│  └──────────────┘     └──────────────┘     └──────┬───────┘                        │
│                                                    │                               │
│                                                    ▼                               │
│                                         ┌────────────────────┐                     │
│                                         │     解析 Action     │                     │
│                                         │                    │                     │
│                                         │  <search> query    │                     │
│                                         │  <answer> answer   │                     │
│                                         └─────────┬──────────┘                     │
│                                                   │                                │
│                                ┌──────────────────┴──────────────────┐             │
│                                │                                     │             │
│                                ▼                                     ▼             │
│                      ┌────────────────┐                  ┌────────────────┐        │
│                      │ action=search  │                  │ action=answer  │        │
│                      │                │                  │                │        │
│                      │  调用检索服务   │                  │   标记完成      │         │
│                      └───────┬────────┘                  └────────────────┘        │
│                              │                                                     │
│                              ▼                                                     │
│                    ┌──────────────────────┐                                        │
│                    │ 拼接检索结果到       │                                          │
│                    │ rolling state        │                                        │
│                    └──────────┬───────────┘                                        │
│                               │                                                    │
│                               ▼                                                    │
│                           下一轮循环                                                │
│                                                                                    │
└────────────────────────────────────────────────────────────────────────────────────┘

for step in range(max_turns):
    if not active_mask.sum():
        break  # 所有样本都已完成

    # Step 1: 裁剪 rolling state 到有效长度
    rollings.batch = cut_to_effective_len(rollings.batch)

    # Step 2: 仅对活跃样本进行 LLM 生成
    rollings_active = DataProto({k: v[active_mask] for k, v in rollings.batch})
    gen_output = generate_with_gpu_padding(rollings_active)

    # Step 3: 后处理响应（截断到 </search> 或 </answer>）
    responses_ids, responses_str = postprocess_responses(gen_output.batch['responses'])

    # Step 4: 执行环境交互（解析动作 + 调用搜索）
    next_obs, dones, valid_action, is_search = execute_predictions(responses_str)

    # Step 5: 更新活跃掩码
    active_mask = active_mask * ~dones

    # Step 6: 处理检索结果为 token ids
    next_obs_ids = process_next_obs(next_obs)

    # Step 7: 更新 rolling state（拼接响应 + 检索结果）
    rollings = update_rolling_state(rollings, responses_ids, next_obs_ids)

    # Step 8: 更新右侧累积输出
    original_right_side = update_right_side(original_right_side, responses_ids, next_obs_ids)

sample_0: [0, 0, 0, 0, 5, 8, 3, 7, 2, 6]   有效长度 6
sample_1: [0, 0, 0, 0, 0, 0, 4, 9, 1, 3]   有效长度 4
sample_2: [0, 0, 0, 7, 2, 5, 8, 3, 6, 1]   有效长度 7

sample_0: [0, 5, 8, 3, 7, 2, 6]
sample_1: [0, 0, 4, 9, 1, 3, 3]
sample_2: [7, 2, 5, 8, 3, 6, 1]

if num_gpus <= 1:
    return self.actor_rollout_wg.generate_sequences(active_batch)

batch_size = active_batch.batch['input_ids'].shape[0]
remainder = batch_size % num_gpus
if remainder == 0:
    return self.actor_rollout_wg.generate_sequences(active_batch)

# padding 补充 batch_size
padded_output = self.actor_rollout_wg.generate_sequences(padded_active_batch)
# 清除多余的 padding 部分
return padded_output

padding_size = num_gpus - remainder
padded_batch = {}

for k, v in active_batch.batch.items():
    # Use first sequence as padding template
    pad_sequence = v[0:1].repeat(padding_size, *[1] * (len(v.shape) - 1))
    padded_batch[k] = torch.cat([v, pad_sequence], dim=0)
padded_active_batch = DataProto.from_dict(padded_batch)
for key in padded_active_batch.batch.keys():
    padded_active_batch.batch[key] = padded_active_batch.batch[key].long()
# Generate with padded batch
padded_output = self.actor_rollout_wg.generate_sequences(padded_active_batch)

input_ids.shape = [13, 4096]
attention_mask.shape = [13, 4096]
position_ids.shape = [13, 4096]

trimmed_batch = {k: v[:-padding_size] for k, v in padded_output.batch.items()}
padded_output.batch = trimmed_batch

Search-R1
  -> actor_rollout_wg.generate_sequences()
    -> FSDP ActorRolloutRefWorker.generate_sequences()
      -> self.rollout.generate_sequences()
        -> vLLMRollout.generate_sequences()
          -> self.inference_engine.generate()

self.inference_engine = LLM(
    actor_module,
    tokenizer=tokenizer,
    model_hf_config=model_hf_config,
    tensor_parallel_size=tensor_parallel_size,
    dtype=config.dtype,
    enforce_eager=config.enforce_eager,
    gpu_memory_utilization=config.gpu_memory_utilization,
    skip_tokenizer_init=False,
    max_model_len=config.prompt_length + config.response_length,
    load_format=config.load_format
)

if '</search>' in resp:
    resp = resp.split('</search>')[0] + '</search>'
elif '</answer>' in resp:
    resp = resp.split('</answer>')[0] + '</answer>'

def _passages2string(retrieval_result):
    format_reference = ''
    for idx, doc_item in enumerate(retrieval_result):
        content = doc_item['document']['contents']
        title = content.split("\n")[0]
        text = "\n".join(content.split("\n")[1:])
        format_reference += f"Doc {idx+1}(Title: {title}) {text}\n"
    return format_reference

My previous action is invalid. 
If I want to search, I should put the query between <search> and </search>. 
If I want to give the final answer, I should put the answer between <answer> and </answer>. Let me try again.

new_input_ids = concatenate_with_padding([
    rollings.batch['input_ids'],  # 之前的 rolling state
    cur_responses,                 # 本轮 LLM 响应
    next_obs_ids                   # 本轮检索结果
])

# 裁剪到 max_prompt_length
new_input_ids = new_input_ids[:, -max_prompt_length:]
new_rollings = DataProto.from_dict({
    'input_ids': new_input_ids[:, -max_len:],
    'position_ids': new_position_ids[:, -max_len:],
    'attention_mask': new_attention_mask[:, -max_len:]
})

max_len = min(self.config.max_prompt_length, effective_len)  
new_rollings = DataProto.from_dict({  
    'input_ids': new_input_ids[:, -max_len:],  # 取最右侧，即丢弃最左侧（system prompt + 原始问题）  
    ...  
})

def run_llm_loop(self, gen_batch, initial_input_ids: torch.Tensor) -> Tuple[Dict, Dict]:
	# original_left_side  → 保存初始 prompt 的固定副本（不变）  
	# original_right_side → 累积所有 responses（从右侧增长）
	return self._compose_final_output(original_left_side, original_right_side, meta_info)

def _select_rm_score_fn(data_source):
    if data_source in ['nq', 'triviaqa', 'popqa', 'hotpotqa', '2wikimultihopqa', 'musique', 'bamboogle']:
        return qa_em.compute_score_em
    else:
        raise NotImplementedError

def compute_score_em(solution_str, ground_truth, format_score=0., score=1.):
    answer = extract_solution(solution_str=solution_str)

    if answer is None:
        return 0
    else:
        if em_check(answer, ground_truth['target']):
            return score
        else:
            return format_score

[prompt 中的示例 <answer> Beijing </answer>]
[模型生成的 <answer> final answer </answer>]

prompt_ids = data_item.batch['prompts']
response_ids = data_item.batch['responses']

valid_prompt_length = data_item.batch['attention_mask'][:prompt_length].sum()
valid_response_length = data_item.batch['attention_mask'][prompt_length:].sum()

valid_prompt_ids = prompt_ids[-valid_prompt_length:]
valid_response_ids = response_ids[:valid_response_length]

sequences = torch.cat((valid_prompt_ids, valid_response_ids))
sequences_str = tokenizer.decode(sequences)

reward_tensor.shape == responses.shape

reward_tensor = torch.zeros_like(data.batch['responses'], dtype=torch.float32)
reward_tensor[i, valid_response_length - 1] = score
batch.batch['token_level_scores'] = reward_tensor
# ...
batch.batch['token_level_rewards'] = batch.batch['token_level_scores']

def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_repeat=1):
    # prepare response group
    # TODO: add other ways to estimate advantages
    if adv_estimator == 'gae':
        values = data.batch['values']
        responses = data.batch['responses']
        response_length = responses.size(-1)
        attention_mask = data.batch['attention_mask']
        response_mask = attention_mask[:, -response_length:]
        token_level_rewards = data.batch['token_level_rewards']
        advantages, returns = core_algos.compute_gae_advantage_return(token_level_rewards=token_level_rewards,
                                                                      values=values,
                                                                      eos_mask=response_mask,
                                                                      gamma=gamma,
                                                                      lam=lam)
        data.batch['advantages'] = advantages
        data.batch['returns'] = returns
    elif adv_estimator == 'grpo':
        token_level_rewards = data.batch['token_level_rewards']
        index = data.non_tensor_batch['uid']
        responses = data.batch['responses']
        response_length = responses.size(-1)
        attention_mask = data.batch['attention_mask']
        response_mask = attention_mask[:, -response_length:]
        advantages, returns = core_algos.compute_grpo_outcome_advantage(token_level_rewards=token_level_rewards,
                                                                        eos_mask=response_mask,
                                                                        index=index)
        data.batch['advantages'] = advantages
        data.batch['returns'] = returns
    else:
        raise NotImplementedError
    return data

def compute_policy_loss(old_log_prob, log_prob, advantages, eos_mask, cliprange):
    negative_approx_kl = log_prob - old_log_prob
    ratio = torch.exp(negative_approx_kl)
    ppo_kl = verl_F.masked_mean(-negative_approx_kl, eos_mask)

    pg_losses = -advantages * ratio
    pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - cliprange, 1.0 + cliprange)

    pg_loss = verl_F.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask)
    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask)
    return pg_loss, pg_clipfrac, ppo_kl

train_ppo.sh / train_grpo.sh
  ↓
verl.trainer.main_ppo
  ↓
RayPPOTrainer.fit()
  ↓
LLMGenerationManager.run_llm_loop()
  ↓
actor_rollout_wg.generate_sequences()
  ↓
多轮 search rollout 得到完整 responses
  ↓
actor_rollout_wg.compute_log_prob()
  ↓
ref_policy_wg.compute_ref_log_prob()    可选
  ↓
critic_wg.compute_values()              PPO 需要
  ↓
reward_fn(batch)
  ↓
compute_advantage(...)
  ├── GAE  → PPO
  └── GRPO → group outcome advantage
  ↓
_create_loss_mask()
  ↓
actor_rollout_wg.update_actor()
  ↓
DataParallelPPOActor.update_policy()
  ↓
core_algos.compute_policy_loss()

def compute_gae_advantage_return(token_level_rewards: torch.Tensor, values: torch.Tensor, eos_mask: torch.Tensor, gamma: torch.Tensor, lam: torch.Tensor):
    with torch.no_grad():
        lastgaelam = 0
        advantages_reversed = []
        gen_len = token_level_rewards.shape[-1]

        for t in reversed(range(gen_len)):
            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
            lastgaelam = delta + gamma * lam * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1)

        returns = advantages + values
        advantages = verl_F.masked_whiten(advantages, eos_mask)
    return advantages, returns

def compute_grpo_outcome_advantage(
	token_level_rewards: torch.Tensor,
	eos_mask: torch.Tensor,
	index: torch.Tensor,
	epsilon: float = 1e-6
):
    response_length = token_level_rewards.shape[-1]
    non_zero_mask = (token_level_rewards != 0)
    scores = (token_level_rewards * non_zero_mask).sum(dim=-1)

    id2score = defaultdict(list)
    id2mean = {}
    id2std = {}

    with torch.no_grad():
        bsz = scores.shape[0]
        for i in range(bsz):
            id2score[index[i]].append(scores[i])
        for idx in id2score:
            if len(id2score[idx]) == 1:
                id2mean[idx] = torch.tensor(0.0)
                id2std[idx] = torch.tensor(1.0)
            elif len(id2score[idx]) > 1:
                id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
                id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
            else:
                raise ValueError(f"no score in prompt index: {idx}")
        for i in range(bsz):
            scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
        scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask

    return scores, scores

def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"):
    responses = data.batch["responses"]
    response_length = responses.size(1)
    token_level_scores = data.batch["token_level_scores"]
    batch_size = data.batch.batch_size[0]
    attention_mask = (
        data.batch["info_mask"]
        if "info_mask" in data.batch
        else data.batch["attention_mask"]
    )
    response_mask = attention_mask[:, -response_length:]

    # compute kl between ref_policy and current policy
    if "ref_log_prob" in data.batch.keys():
        kld = core_algos.kl_penalty(
            data.batch["old_log_probs"],
            data.batch["ref_log_prob"],
            kl_penalty=kl_penalty,
        )  # (batch_size, response_length)
        kld = kld * response_mask
        beta = kl_ctrl.value
    else:
        beta = 0
        kld = torch.zeros_like(response_mask, dtype=torch.float32)

    token_level_rewards = token_level_scores - beta * kld

    current_kl = masked_mean(kld, mask=response_mask, axis=-1)  # average over sequence
    current_kl = torch.mean(current_kl, dim=0).item()

    # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837
    kl_ctrl.update(current_kl=current_kl, n_steps=batch_size)
    data.batch["token_level_rewards"] = token_level_rewards

    metrics = {"critic/kl": current_kl, "critic/kl_coeff": beta}

    return data, metrics

actor_rollout_ref.actor.use_kl_loss=true
actor_rollout_ref.actor.kl_loss_coef=0.001
actor_rollout_ref.actor.kl_loss_type=low_var_kl

batch.batch['token_level_rewards'] = batch.batch['token_level_scores']

kld = core_algos.kl_penalty(
    logprob=log_prob,
    ref_logprob=ref_log_prob,
    kl_penalty=self.config.kl_loss_type,
)
kl_loss = masked_mean(kld, response_mask)
policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef

for epoch in range(self.config.trainer.total_epochs):
    for batch_dict in self.train_dataloader:
        batch: DataProto = DataProto.from_single_dict(batch_dict)
        batch = batch.repeat(
            repeat_times=self.config.actor_rollout_ref.rollout.n_agent,
            interleave=True,
        )

        gen_batch = batch.pop(
            batch_keys=["input_ids", "attention_mask", "position_ids"]
        )

gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)
batch.non_tensor_batch["uid"] = np.array(
    [str(uuid.uuid4()) for _ in range(len(batch.batch))],
    dtype=object,
)
batch = batch.repeat(
    repeat_times=self.config.actor_rollout_ref.rollout.n,
    interleave=True,
)
batch = batch.union(gen_batch_output)

@register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
def generate_sequences(self, prompts: DataProto):
    prompts = prompts.to('cuda')
    recompute_log_prob = prompts.meta_info.get('recompute_log_prob', True)

    assert self._is_rollout

    if self._is_offload_param:
        load_fsdp_param_and_grad(
            module=self.actor_module_fsdp,
            device_id=torch.cuda.current_device(),
            load_grad=self._is_offload_grad,
        )

    prompts.batch = prompts.batch.cuda()
    prompts.meta_info.update({
        'eos_token_id': self.tokenizer.eos_token_id,
        'pad_token_id': self.tokenizer.pad_token_id,
    })

    with self.rollout_sharding_manager:
        prompts = self.rollout_sharding_manager.preprocess_data(prompts)
        output = self.rollout.generate_sequences(prompts=prompts)
        output = self.rollout_sharding_manager.postprocess_data(output)

    if self._is_actor and recompute_log_prob:
        old_log_probs = self.actor.compute_log_prob(data=output)
        output.batch['old_log_probs'] = old_log_probs

    output = output.to('cpu')

    if self._is_offload_param:
        offload_fsdp_param_and_grad(...)

    torch.cuda.empty_cache()
    return output

llamafactory-cli train examples/train_lora/qwen3_lora_sft.yaml
llamafactory-cli export examples/merge_lora/qwen3_lora_sft.yaml

### model
model_name_or_path: Qwen/Qwen3-4B-Instruct-2507
trust_remote_code: true

### method
stage: sft
do_train: true
finetuning_type: lora
lora_rank: 8
lora_target: all

### dataset
dataset: identity,alpaca_en_demo
template: qwen3_nothink
cutoff_len: 2048
max_samples: 1000
preprocessing_num_workers: 16
dataloader_num_workers: 4

### output
output_dir: saves/qwen3-4b/lora/sft
logging_steps: 10
save_steps: 500
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]

### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 8
learning_rate: 1.0e-4
num_train_epochs: 3.0
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null

### eval
# eval_dataset: alpaca_en_demo
# val_size: 0.1
# per_device_eval_batch_size: 1
# eval_strategy: steps
# eval_steps: 500

# alpaca
[
    {
        "instruction": "任务指令",
        "input": "可选的输入上下文",
        "output": "期望的输出响应"
    }
]
# sharegpt
[
    {
        "conversations": [
            {
                "from": "human",
                "value": "用户说的话"
            },
            {
                "from": "gpt",
                "value": "助手的回复"
            },
            {
                "from": "human",
                "value": "用户下一句话"
            }
        ],
        "system": "可选的系统提示词"
    }
]

{
  "my_dataset": {
    "file_name": "my_data.json"
  }
}

{
  "chat_dataset": {
    "file_name": "chat.json",
    "formatting": "sharegpt",
    "columns": {
      "messages": "conversations"
    }
  }
}

{
  "openai_dataset": {
    "file_name": "openai.json",
    "formatting": "sharegpt",
    "columns": {
      "messages": "messages"
    },
    "tags": {
      "role_tag": "role",
      "content_tag": "content",
      "user_tag": "user",
      "assistant_tag": "assistant",
      "system_tag": "system"
    }
  }
}

original_left_side = {
    'input_ids': initial_input_ids[:, -self.config.max_start_length:]
}

original_right_side = {
    'responses': initial_input_ids[:, []],
    'responses_with_info_mask': initial_input_ids[:, []]
}

    def _compose_final_output(self, left_side: Dict, right_side: Dict, meta_info: Dict) -> Tuple[Dict, Dict]:
        """Compose final generation output."""
        final_output = right_side.copy()
        final_output['prompts'] = left_side['input_ids']
        
        # Combine input IDs
        final_output['input_ids'] = torch.cat([
            left_side['input_ids'],
            right_side['responses']
        ], dim=1)
        
        # Create attention mask and position ids
        final_output['attention_mask'] = torch.cat([
            self.tensor_fn.create_attention_mask(left_side['input_ids']),
            self.tensor_fn.create_attention_mask(final_output['responses'])
        ], dim=1)
        final_output['info_mask'] = torch.cat([
            self.tensor_fn.create_attention_mask(left_side['input_ids']),
            self.tensor_fn.create_attention_mask(final_output['responses_with_info_mask'])
        ], dim=1)
        
        final_output['position_ids'] = self.tensor_fn.create_position_ids(
            final_output['attention_mask']
        )
        
        final_output = DataProto.from_dict(final_output)
        final_output.meta_info.update(meta_info)
        
        return final_output

    def _create_loss_mask(self, batch, metrics):
        """Create loss mask for state tokens."""
        response_length = batch.batch["responses"].shape[-1]
        response_mask = batch.batch["attention_mask"][:, -response_length:]

        loss_mask = batch.batch["info_mask"][:, -response_length:]
        batch.batch["loss_mask"] = loss_mask

        metrics.update(
            {
                "state_tokens/total": loss_mask.sum().item(),
                "state_tokens/coverage": (loss_mask.sum() / response_mask.sum()).item(),
            }
        )

        return batch, metrics

def _compute_loss(self, batch):
    loss_mask = batch.pop('loss_mask')[:, :-1].reshape(-1).cuda()
    labels = batch['input_ids'][:, 1:].cuda()
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        output = self.fsdp_model(input_ids=batch['input_ids'],
                                 attention_mask=batch['attention_mask'],
                                 position_ids=batch['position_ids'],
                                 use_cache=False)  # prevent model thinks it it generating
    logits = output.logits
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels.contiguous()
    # Flatten the tokens
    loss_fct = nn.CrossEntropyLoss(reduction='none')
    shift_logits = shift_logits.view(-1, self.model.config.vocab_size)
    shift_labels = shift_labels.view(-1)
    # Enable model parallelism
    shift_labels = shift_labels.to(shift_logits.device)
    loss = loss_fct(shift_logits, shift_labels)
    loss = loss * loss_mask

max_prompt_length = max_start_length + max_response_length * (max_turns - 1) + max_obs_length * max_turns

if self.val_reward_fn is not None and self.config.trainer.get(
    "val_before_train", True
):
    val_metrics = self._validate()
    pprint(f"Initial validation metrics: {val_metrics}")
    logger.log(data=val_metrics, step=self.global_steps)
    if self.config.trainer.get("val_only", False):
        return