Support preformatted datasets

[shimizust]

Motivation

• Generic datasets like sharegpt/ultrachat do not have user/assistant tokens, and we expect to have to apply a specific model's chat template to it. However, there are use cases where we want to train a draft model on the dataset used during fine-tuning of the target model which has its own system prompt and may have been saved as the raw prompt, along with raw target model generations (including user/assistant/eot tokens).
• Currently, we would need to reverse-engineer the chat template or create new templates with custom system prompts, which is inconvenient and error-prone.

Modifications

• Add support for passing pre-formatted conversational text directly without applying additional chat templating by passing in the --is-preformatted flag to the training scripts.
• If --is-preformatted is passed, then we expect json data in this format:

{
  "id": ...,
  "text": "<|im_start|>system\nYou are a ...<|im_end|>\n<|im_start|>user\nHello...<|im_start|>assistant\n<think>..."
}

The current dataset format for conversations that are not pre-formatted:

{
  "id": ...,
  "conversations": [{"role": "user", "content": ...}, {"role": "assistant", "content": ...}]
}

Accuracy Test

• Added visual debugging test for preprocessing conversations with/without pre-formatting.

• Ran training successfully without preformatting (sharegpt) and with pre-formatted dataset:

torchrun \
    --nnodes 1 \
    --nproc_per_node 2 \
    train_eagle3_online.py \
    --target-model-path <TARGET_MODEL_PATH> \
    --draft-model-config <MODEL_CONFIG_PATH> \
    --train-data-path <PATH_TO_PREFORMATTED_DATASET> \
    --output-dir <OUTPUT_DIR> \
    --num-epochs 10 \
    --batch-size 1 \
    --learning-rate 1e-4 \
    --max-length 8192 \
    --tp-size 1 \
    --chat-template qwen \
    --is-preformatted \
    --cache-dir /dev/shm/cache \
    --log-steps 50 \
    --report-to mlflow

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[FlamingoPg]

LGTM, any other comment? @FrankLeeeee @shuaills

[shimizust]

Thanks @FlamingoPg , lmk your thoughts @FrankLeeeee @shuaills