[iamtatsuki05]

Summary

Adds document-boundary-aware sequence packing that both run_mlm.py (encoders) and run_clm.py (decoders) can switch on with the same config flags. Default behaviour is unchanged (packing: false).

• Packing: trl.data_utils.pack_dataset(strategy="bfd") packs whole documents into rows of packing_seq_length tokens (documents longer than a row are truncated); every row carries seq_lengths. packing_strategy: wrapped reproduces the old concatenate-and-chunk (group_texts) behaviour.
• Decoders (Llama / Qwen / Mistral ...): PackedCausalLMCollator emits per-document position_ids, labels (padding and, optionally, document-start tokens set to -100), the FlashAttention varlen kwargs (cu_seq_lens_q/k, max_length_q/k) and no attention_mask. flash_attention_2 consumes the varlen kwargs; sdpa / eager get the block-diagonal mask that transformers (>= 4.53, torch >= 2.6) derives from the position_ids restarts, which requires use_cache=False (set by the script).
• ModernBERT: packing_encoder_mode=unpad splits rows back into documents; the model's own unpadding re-packs them with cu_seqlens (true packing with flash_attention_2).
• BERT / RoBERTa-style encoders: packing_encoder_mode=mask keeps the packed row and passes a 3D block-diagonal attention_mask plus per-document position_ids.
• Library versions that would silently drop the boundaries are rejected at start-up; packing_seq_length > max_seq_length, packing + pad_to_max_length and mask mode for ModernBERT are rejected as well.

New config keys: packing, packing_strategy, packing_seq_length, packing_encoder_mode (MLM), packing_mask_document_starts (CLM). See scripts/constract_llm/train/pt/README.md / README_JA.md.

Verification

• Local tox-equivalent (Python 3.12, transformers 4.56.1, trl 0.22.2, torch 2.8 CPU): pytest tests/ -n auto 146 passed, ruff check / ruff format --check clean.
• New tests (tests/mirei/constract_llm/train/language_model/test_packing.py, 16): a packed Llama batch (2 rows + padding segment) matches per-document logits for eager and sdpa; a packed BERT batch with the 3D mask matches per-document logits; collator shapes / positions / labels / varlen kwargs; unpad-mode masking; argument validation.
• CPU end-to-end smoke of run_mlm.py (mask and unpad modes) and run_clm.py with packing: true on a tiny random model: train and eval complete, the packing_seq_length guard raises as intended.
• Not verified on GPU: the flash_attention_2 varlen path (decoder) and ModernBERT unpadding under packing. The kwargs follow transformers.modeling_flash_attention_utils._flash_attention_forward (the varlen branch reshapes (batch, seq) row-major); a packed-vs-per-document logits check on a GPU is recommended before a production run.

Notes

• With packing on, a "batch" is a number of packed rows, i.e. about per_device_train_batch_size × packing_seq_length tokens per micro-batch; step counts and the meaning of the total batch change accordingly.
• use_cache=False is persisted in the saved config.json of packed CLM runs; re-enable it for generation.
• This PR is based on main and touches only the tokenize → group / collator hunks of run_mlm.py / run_clm.py.