EXTREMELY SLOW (unusable) towards end of tokenization of dataset with long multi turn conversations

Author: Nero10578Created Mar 7, 2025Updated Jul 18, 2026

Labelsbug

Please check that this issue hasn't been reported before.

I searched previous Bug Reports didn't find any similar reports.

Expected Behavior

Expected tokenization to work extremely fast as in the old commits in from this commit and older 339f3c67e2d6855340b5958274ea539517829baa

Current behaviour

Ever since a change to how the chat_templates works back in this commit 10cfecf02e8829de749708c2588dc76be3a156d2 the tokenization has become literally unusable for my dataset that contains long multi turn conversation. The speed seems to become exponentially slower towards the end of the tokenization since it uses less and less CPU cores as it finishes most of the examples.

Since that commit I have just been replacing the chat_template.py with the older one and it worked fine, but now since this commit b194e17c2847348bb102d24e86775dfb77c9c5fe this hacky fix stopped working as the new chat_template.py is much different. So now newer axolotl version are unusable to me.

I can help the performance to be better by changing num_proc in datasets.py to a huge number but it still eventually slows to a crawl at the end.

Steps to reproduce

Run preprocess on using the latest axolotl commit on a long multi turn conversation dataset.

Config yaml

yaml

base_model: /home/owen/models/Mistral-Small-24B-Instruct-2501
tokenizer_type: AutoTokenizer

train_on_inputs: false
group_by_length: false
load_in_8bit: false
load_in_4bit: true
strict: false

sequence_len: 8192
bf16: auto
tf32: false
flash_attention: true

shuffle_merged_datasets: true

# lora_on_cpu: true
# gpu_memory_limit: 20GiB

#Data
datasets:
  - path: /home/owen/datasets/RPMax-v1.4-formatted/conversations-escaped-formatted_split_1.json
    type: chat_template
    field_messages: conversations
    message_field_role: from
    message_field_content: value


warmup_steps: 20
dataset_prepared_path: ./lora_last_run_prepared

# Iterations
num_epochs: 1
saves_per_epoch: 8
saves_total_limit: 8

# Evaluation
val_set_size: 0.0025
eval_max_new_tokens: 128
eval_sample_packing: false
evals_per_epoch: 8

# LoRA
output_dir: ./qlora_out_rex
adapter: qlora
lora_model_dir:
lora_r: 128
lora_alpha: 128
lora_target_linear: true

save_safetensors: true

peft_use_rslora: false
loraplus_lr_ratio: 16

# Sampling
sample_packing: true
pad_to_sequence_len: true

# Batching
gradient_accumulation_steps: 16
micro_batch_size: 1
gradient_checkpointing: false
gradient_checkpointing_kwargs:
  use_reentrant: false

# wandb
wandb_mode: # "offline" to save run metadata locally and not sync to the server, "disabled" to turn off wandb
wandb_project: 
wandb_entity: # A wandb Team name if using a Team
wandb_watch:
wandb_name:
wandb_run_id: # Set the ID of your wandb run
wandb_log_model: # "checkpoint" to log model to wandb Artifacts every `save_steps` or "end" to log only at the end of training

# Optimizer
optimizer: adamw_torch
lr_scheduler: cosine
learning_rate: 0.00001

# Misc
auto_resume_from_checkpoints: true
logging_steps: 1
weight_decay: 0.0

special_tokens:
  pad_token: <pad>

plugins:
  - axolotl.integrations.liger.LigerPlugin
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
cut_cross_entropy: true
liger_rope: true
liger_rms_norm: true
liger_layer_norm: true
liger_glu_activation: true
#liger_fused_linear_cross_entropy: true

# Multi-GPU
# deepspeed: /home/owen/axolotl/deepspeed_configs/zero3_bf16_cpuoffload_params.json
fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_limit_all_gathers: true
  fsdp_sync_module_states: true
  fsdp_offload_params: false
  fsdp_use_orig_params: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: MistralDecoderLayer
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_activation_checkpointing: true

Possible solution

I think what chat_template.py does is find start and end points for each turn by checking the raw tokens? Shouldn't it just use the json structure instead when it can, where there is already a clear conversation start and end point for each turn.

Which Operating Systems are you using?

Linux
macOS
Windows

Python Version

3.12

axolotl branch-commit

main/16dc6ee68d05bfbef607a8ee06e00fca1d8c4739

Acknowledgements

My issue title is concise, descriptive, and in title casing.
I have searched the existing issues to make sure this bug has not been reported yet.
I am using the latest version of axolotl.
I have provided enough information for the maintainers to reproduce and diagnose the issue.

Source: axolotl-ai-cloud/axolotl

View original on GitHub View discussion on GitHub