Question about Inference Time during Training

[jingfanke]

Thank you for your excellent work and for sharing this impressive project!

I encountered some efficiency issues during training and would like to ask for your advice. Specifically, the inference step seems to be extremely slow in my setup. Here's the configuration I'm using:

• Dataset: MuSiQue
• batch_size: 16
• retrieval_batch_size: 32
• Hardware: 8 × A100 (40GB)
• Retriever model loaded on 1 GPU
• Two Llama-3.1–8B-Instruct models loaded across the remaining 7 GPUs using data parallelism for inference

Under this setup, computing the PPL scores for one batch (16 × 32 = 512 examples) takes over 2 minutes, which becomes a significant bottleneck in training.

I wonder if you have any suggestions or recommended optimizations to reduce the inference time?

Any guidance would be greatly appreciated. Thank you again for your contributions!

[leeds1219]

Thanks for the detailed feedback!

In my setup, I only used 1–2 GPUs, so I didn’t implement batch processing in the _score function — it's currently running one example at a time, which explains the slow PPL computation. I was planning to optimize that later, but haven’t gotten around to it yet.

If you want to implement batching yourself in the meantime, it should be straightforward — just make sure to handle padding properly when computing the loss.

e.g.

pad_token_id = tokenizer.pad_token_id
loss_fct = torch.nn.CrossEntropyLoss(ignore_index=pad_token_id, reduction='none')

batch_size = 2
all_nlls = []

for i in range(0, len(input_prompts), batch_size):
    batch_prompt_indices = prompt_indices[i:i+batch_size]
    batch_input_prompts = input_prompts[i:i+batch_size]
    batch_output_prompts = output_prompts[i:i+batch_size]
    break  # Process only the first batch

input_ids = []
attention_masks = []

# 1) Calculate max_len
max_len = max(
    len(tokenizer.encode(input_prompts[idx], add_special_tokens=False) +
        tokenizer.encode(output_prompts[idx], add_special_tokens=False))
    for idx in range(len(batch_prompt_indices))
) + 1

# 2) Create input_ids and attention_mask (right padding)
for idx in range(len(batch_prompt_indices)):
    inp = batch_input_prompts[idx]
    out = batch_output_prompts[idx]

    encoded_inp = tokenizer.encode(inp, add_special_tokens=False)
    encoded_out = tokenizer.encode(out, add_special_tokens=False)

    full_seq = encoded_inp + encoded_out

    padding_len = max_len - len(full_seq)

    padded = full_seq + [pad_token_id] * padding_len  # Right padding
    input_ids.append(padded)

    attention_masks.append([1] * len(full_seq) + [0] * padding_len)

input_ids = torch.tensor(input_ids).to(device)
attention_masks = torch.tensor(attention_masks).to(device)

# 3) Create shift_labels (based on right padding)
shift_labels = torch.full_like(input_ids, pad_token_id)

for i in range(input_ids.size(0)):
    tokens = input_ids[i]
    padding_mask = (tokens == pad_token_id)

    seq_len = (~padding_mask).sum().item()
    real_tokens = tokens[:seq_len]
    shifted = torch.cat([real_tokens[1:], torch.tensor([pad_token_id], device=device)])

    shift_labels[i, :seq_len] = shifted

# 4) Model output
outputs = llm_model(input_ids=input_ids, attention_mask=attention_masks)
logits = outputs.logits

# 5) Shift logits and labels (shift_labels is already shifted)
shift_logits = logits[:, :-1, :].contiguous()
shift_labels_cut = shift_labels[:, :-1].contiguous()

# 6) Compute token-wise loss
loss = loss_fct(
    shift_logits.view(-1, logits.size(-1)),
    shift_labels_cut.view(-1)
).view(shift_labels_cut.size())

# 7) Create loss mask
loss_mask = (shift_labels_cut != pad_token_id).float()

# 8) Total loss and token count per sample
sum_loss_per_sample = (loss * loss_mask).sum(dim=1)           # Sum of token losses
num_tokens_per_sample = loss_mask.sum(dim=1).clamp(min=1)     # Number of tokens

# 9) Average negative log likelihood (NLL) per sample
avg_nll_per_sample = sum_loss_per_sample / num_tokens_per_sample

print("Batch negative log probability per sample:", avg_nll_per_sample)
print(sum_loss_per_sample)

Output should be some thing like

Batch negative log probability per sample: tensor([2.9080, 2.9204], device='cuda:0', grad_fn=<DivBackward0>)
tensor([782.2537, 779.7516], device='cuda:0', grad_fn=<SumBackward1>)

Which should be close to the original implementation

individual_nlls = []

for i in range(0, len(input_prompts), batch_size):
    batch_prompt_indices = prompt_indices[i:i+batch_size]
    break

for inp, out in zip(input_prompts[i:i+batch_size], output_prompts[i:i+batch_size]):
    encoded_inp = tokenizer.encode(inp, add_special_tokens=False)
    encoded_out = tokenizer.encode(out, add_special_tokens=False)

    full_seq = encoded_inp + encoded_out
    input_ids_single = torch.tensor([full_seq]).to(device)
    attention_mask_single = (input_ids_single != pad_token_id).long().to(device)

    outputs_single = llm_model(input_ids=input_ids_single, attention_mask=attention_mask_single)
    logits_single = outputs_single.logits

    shift_logits_single = logits_single[:, :-1, :].contiguous()
    shift_labels_single = input_ids_single[:, 1:].contiguous()

    loss_single = loss_fct(shift_logits_single.view(-1, logits_single.size(-1)), shift_labels_single.view(-1))
    loss_single = loss_single.view(shift_labels_single.size())

    loss_mask_single = (shift_labels_single != pad_token_id).float()
    neg_log_prob_single = (loss_single * loss_mask_single).sum().item()
    
    individual_nlls.append(neg_log_prob_single)

print("Individual neg log probs:", individual_nlls)

Output should be some thing like

Individual neg log probs: [782.2535400390625, 779.75146484375]