ChipYTY/titans_NPC

1## BABILong（QA1 / 32k）数据处理与训练数据流说明2 3本文档描述当前仓库里 **BABILong QA1（32k.json）** 在训练脚本中的实际处理方式：从原始 JSON，到 tokenizer、padding、labels、DataLoader，再到喂给 `QwenTitansForBABILong` 的整条数据流。4 5代码入口：6 7- `examples/train_qwen_titans_babilong.py`8 9---10 11## 数据源与样本格式12 13默认数据路径（可改）：14 15- `TrainingConfig.data_path = /data/yty/BABILong/babilong-train-5k-samples/data/qa1/32k.json`16 17文件内容为一个 JSON 列表，每条样本大体包含：18 19- `input`：长上下文（故事/事实）20- `question`：问题21- `target`：答案（短文本）22 23训练脚本会把它拼成 prompt：24 25```26{input}27 28Question: {question}29Answer:30```31 32并把答案拼接为（答案前加空格）：33 34```35 {target}36```37 38---39 40## 关键目标：固定长度样本（FSDP/DDP 必须）41 42当前实现 **强制每条样本输出固定长度 `config.max_length`（默认 32768）**，原因：43 44- 模型前向会按 `chunk_size` 把序列分 chunk 循环处理45- 在 FSDP/DDP 下，如果不同 rank 的序列长度不同 → chunk 数不同 → collectives 顺序不一致 → NCCL watchdog 超时46 47因此数据侧必须固定长度，保证每步每个 rank 的 chunk 次数一致。48 49相关参数：50 51- `TrainingConfig.max_length`：固定输出长度（默认 32768）52- `TrainingConfig.answer_reserve_tokens`：给答案预留 token 数（默认 64）53 54---55 56## Dataset：`BABILongDataset.__getitem__` 的处理流程57 58位置：`examples/train_qwen_titans_babilong.py`59 60### Step 1：tokenize prompt（截断）61 62- 对 prompt 进行 tokenize63- 最大长度限制为 `max_length - answer_reserve_tokens`64- `add_special_tokens=True`（让 tokenizer 自己加 BOS/EOS 等需要的特殊 token）65 66### Step 2：tokenize answer（不加特殊 token）67 68- 对 `" {target}"` tokenize69- `add_special_tokens=False`70 71### Step 3：拼接并截断到 `max_length`72 73- 先算 prompt token 数 `len(prompt_ids)`74- answer 只保留剩余可用空间：`available = max_length - len(prompt_ids)`75- `input_ids = concat(prompt_ids, answer_ids[:available])`76 77### Step 4：构造 `labels`（只监督答案）78 79- `labels` 初始全为 `-100`80- 只有答案 token 的位置才写入对应 token id81- 这样 loss 只在答案 token 上计算（prompt 与 padding 不参与 loss）82 83### Step 5：padding 到固定长度 + attention_mask84 85如果拼接后长度 `< max_length`：86 87- `input_ids` 右侧 pad 到 `max_length`（pad_id = tokenizer.pad_token_id）88- `labels` pad 的部分保持 `-100`89- `attention_mask`：90  - 真 token 为 191  - padding 为 092 93> 备注：脚本在 `main()` 里如果发现 `tokenizer.pad_token is None`，会设置 `pad_token = eos_token`，确保有 pad_id。94 95---96 97## DataLoader 与分布式采样98 99### DataLoader100 101- `batch_size = 1`（32k 序列 + chunk streaming，一般只能 1）102- `collate_fn` 只做 stack（Dataset 已固定长度，不做动态 padding）103- `num_workers = 0`（避免多进程复制大张量带来的额外开销/不稳定）104 105### 训练/验证切分106 107- `random_split(full_dataset, [train_size, eval_size], generator=manual_seed(config.seed))`108- 默认 `train_ratio=0.9`109 110### 分布式（torchrun）111 112当使用 `torchrun` 启动时：113 114- 训练集：`DistributedSampler(..., shuffle=True, seed=config.seed)`115- 验证集：`DistributedSampler(..., shuffle=False)`116- 每个 epoch 会调用 `train_sampler.set_epoch(epoch)`，保证各 rank shuffle 一致117 118---119 120## 喂给模型的数据张量形状121 122由于固定长度：123 124- `input_ids`: `[B, max_length]`（默认 `[1, 32768]`）125- `attention_mask`: `[B, max_length]`126- `labels`: `[B, max_length]`127 128模型内部再按 `chunk_size`（默认 4096）切成 8 个 chunk 进行 streaming。129 130---131 132## 训练与日志（跟数据流相关的行为）133 134- **梯度累积**：`gradient_accumulation_steps=8`135  - 每 8 个 micro-batch 才做一次 optimizer step136- **每 80 个 batch 输出一次**：137  - `--log_every_batches 80`（默认 80）138  - 会自动换算成 `logging_steps = ceil(log_every_batches / gradient_accumulation_steps)`139  - 并在 rank0 额外 `logger.info(...)` 打一行到终端，方便 `tee` 保存140 141---142 143## 运行方式（推荐）144 145### 8 卡 + FSDP146 147```bash148CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \149torchrun --standalone --nproc_per_node=8 \150  examples/train_qwen_titans_babilong.py --fsdp --log_every_batches 80151```152 153### 快速小跑（2 卡调试）154 155```bash156CUDA_VISIBLE_DEVICES=0,1 \157torchrun --standalone --nproc_per_node=2 \158  examples/train_qwen_titans_babilong.py --fsdp --max_samples 8 --num_epochs 1 --eval_steps 1000000159```160 161---162 163## 训练产物（输出）164 165默认输出目录：166 167- `TrainingConfig.output_dir = ./outputs/qwen_titans_babilong`168 169默认只保存一个 final checkpoint（覆盖写入）：170 171- `final_memory_checkpoint.pt`172 173内容包括：174 175- `memory_state_dict`：只包含 `long_term_memory` / `memory_gate` 的参数（体积更小）176- `global_step`177 178