I am pretraining a Llama-based causal language model from scratch on a custom tokenized binary dataset (where data is stored as raw integer IDs).
However, during inference, the generated sequence consistently skips the first expected token and starts from a later offset. For instance, the generated token id sequence is [2,3,4...] when [1,2,3,4...] is expected.

Here is my code:

### load_data.py

import torch
import numpy as np
from torch.utils.data import IterableDataset
import glob
import os
import random

class BinaryDataset(IterableDataset):
    def __init__(self, data_dir, block_size, dtype=np.uint16):
        self.block_size = block_size
        self.dtype = dtype
        self.file_paths = sorted(glob.glob(os.path.join(data_dir, "*.bin")))

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            files = self.file_paths
        else:
            per_worker = int(np.ceil(len(self.file_paths) / float(worker_info.num_workers)))
            worker_id = worker_info.id
            files = self.file_paths[worker_id * per_worker : (worker_id + 1) * per_worker]
        random.shuffle(files)

        for file_path in files:
            data = np.fromfile(file_path, dtype=self.dtype)
            if len(data) - self.block_size <= 0:
                continue
                
            for idx in range(0, len(data) - self.block_size, self.block_size):
                chunk = data[idx : idx + self.block_size + 1].astype(np.int64)
                yield {
                    "input_ids": torch.from_numpy(chunk[:-1]),
                    "label_ids": torch.from_numpy(chunk[1:])
                }
train_dataset = BinaryDataset("./data/train", block_size=4096)
eval_dataset = BinaryDataset("./data/eval", block_size=4096)


### model.py
from transformers import LlamaConfig, LlamaForCausalLM

config = LlamaConfig(
    bos_token_id=1,                
    eos_token_id=2,               
    vocab_size=4096,     
    hidden_size=2048,             
    intermediate_size=5632,        
    num_attention_heads=16,        
    num_key_value_heads=4,         
    num_hidden_layers=32,
    max_position_embeddings=4096,
    tie_word_embeddings=True,
    attn_implementation="flash_attention_2"
)

model = LlamaForCausalLM(config)

### train.py
from model import model
from load_data import train_dataset, eval_dataset
from transformers import TrainingArguments, Trainer, LlamaForCausalLM

training_args = TrainingArguments(
    output_dir="./output",
    #num_train_epochs=1,
    per_device_train_batch_size=9, 
    max_steps=65000, 
    warmup_steps=0.02, 
    lr_scheduler_type="cosine",

    bf16=True,                     
    tf32=True,                     
    bf16_full_eval=False,
    
    use_liger_kernel = True,          
    dataloader_num_workers=8,      
    dataloader_pin_memory=True,   
    
    logging_steps=500,
    save_steps=5000,
    save_total_limit=100,

    eval_strategy = "steps",
    eval_steps = 5000,          
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

### inference.py
import torch
from transformers import AutoModelForCausalLM

device = "cuda" if torch.cuda.is_available() else "cpu"
model_path = "./output/checkpoint-60000" 
model = AutoModelForCausalLM.from_pretrained(model_path).to(device)

#expected:[1,3,549,548,25,3,550,548,31,8,549,550,3,551,548,25,5,550,549,10,548,550,14,549,551,548]
raw_ids = [1,3,549,548,25,3,550,548,31,8,549,550,3,551,548,25,5,550]
input_ids = torch.tensor([raw_ids]).to(device)
attention_mask = torch.ones_like(input_ids)

outputs = model.generate(
    bos_token_id=1,               
    eos_token_id=2,               
    input_ids=input_ids,
    attention_mask=attention_mask,
    max_new_tokens=20,     
)


generated_ids = outputs[0].tolist()
print("generated Token IDs:", generated_ids[len(raw_ids):])
```