Add example/ folder with training scripts

example/README.md

@@ -0,0 +1,131 @@
+# BitCPM4 Continue Pretrain Example
+
+This project provides scripts for continue pretraining **BitCPM4-CANN-1B-unquantized**.
+
+## Environment Setup
+
+### Docker Image
+
+Use the following Huawei NPU image:
+
+```
+swr.cn-south-1.myhuaweicloud.com/ascendhub/mindspeed-llm:openeuler22.03-mindspeed-llm-2.3.0-a3-arm
+```
+
+Other Huawei NPU images may also work but have not been fully tested.
+
+### Install Dependencies
+
+After entering the container, install the Python dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Dependency list:
+
+| Package | Version |
+| --- | --- |
+| transformers | 4.46.3 |
+| tokenizers | 0.20.3 |
+| accelerate | 1.1.1 |
+| deepspeed | 0.16.2 |
+| datasets | 3.1.0 |
+| safetensors | 0.4.5 |
+| pyarrow | 17.0.0 |
+| tensorboard | 2.18.0 |
+
+## Dataset
+
+The test dataset used is [C4-Pro](https://huggingface.co/datasets/gair-prox/c4-pro), stored in parquet format after downloading.
+
+## Usage
+
+Modify the path configuration in `run.sh`:
+
+```bash
+MODEL_PATH="/path/to/BitCPM4-CANN-1B-unquantized/"
+DATA_PATH="/path/to/c4-pro/data/your_file.parquet"
+```
+
+Then start training:
+
+```bash
+bash run.sh
+```
+
+By default, the script trains for 500 steps using 8 devices, DeepSpeed ZeRO-2, and bf16 precision.
+
+## Training Results Reference
+
+Below is the loss curve for the first 100 steps (learning rate warmup covers the first 50 steps):
+
+| Step | Loss | Learning Rate | Epoch |
+| --- | --- | --- | --- |
+| 2 | 2.7920 | 1.60e-06 | 0.01 |
+| 4 | 2.8012 | 3.20e-06 | 0.02 |
+| 6 | 2.7984 | 4.80e-06 | 0.03 |
+| 8 | 2.7839 | 6.40e-06 | 0.04 |
+| 10 | 2.8084 | 8.00e-06 | 0.05 |
+| 12 | 2.8064 | 9.60e-06 | 0.06 |
+| 14 | 2.7994 | 1.12e-05 | 0.07 |
+| 16 | 2.7463 | 1.28e-05 | 0.08 |
+| 18 | 2.7580 | 1.44e-05 | 0.09 |
+| 20 | 2.8007 | 1.60e-05 | 0.10 |
+| 22 | 2.8916 | 1.76e-05 | 0.12 |
+| 24 | 2.8144 | 1.92e-05 | 0.13 |
+| 26 | 2.7723 | 2.08e-05 | 0.14 |
+| 28 | 2.7556 | 2.24e-05 | 0.15 |
+| 30 | 2.7414 | 2.40e-05 | 0.16 |
+| 32 | 2.7469 | 2.56e-05 | 0.17 |
+| 34 | 2.7428 | 2.72e-05 | 0.18 |
+| 36 | 2.7392 | 2.88e-05 | 0.19 |
+| 38 | 2.7132 | 3.04e-05 | 0.20 |
+| 40 | 2.7008 | 3.20e-05 | 0.21 |
+| 42 | 2.7547 | 3.36e-05 | 0.22 |
+| 44 | 2.7151 | 3.52e-05 | 0.23 |
+| 46 | 2.7119 | 3.68e-05 | 0.24 |
+| 48 | 2.7029 | 3.84e-05 | 0.25 |
+| 50 | 2.6803 | 4.00e-05 | 0.26 |
+| 52 | 2.6980 | 4.00e-05 | 0.27 |
+| 54 | 2.6923 | 4.00e-05 | 0.28 |
+| 56 | 2.7068 | 4.00e-05 | 0.29 |
+| 58 | 2.6965 | 4.00e-05 | 0.30 |
+| 60 | 2.7179 | 3.99e-05 | 0.31 |
+| 62 | 2.7119 | 3.99e-05 | 0.32 |
+| 64 | 2.7178 | 3.99e-05 | 0.33 |
+| 66 | 2.7069 | 3.99e-05 | 0.35 |
+| 68 | 2.6870 | 3.98e-05 | 0.36 |
+| 70 | 2.6775 | 3.98e-05 | 0.37 |
+| 72 | 2.7038 | 3.98e-05 | 0.38 |
+| 74 | 2.6924 | 3.97e-05 | 0.39 |
+| 76 | 2.7061 | 3.97e-05 | 0.40 |
+| 78 | 2.6929 | 3.96e-05 | 0.41 |
+| 80 | 2.6787 | 3.96e-05 | 0.42 |
+| 82 | 2.6749 | 3.95e-05 | 0.43 |
+| 84 | 2.6909 | 3.94e-05 | 0.44 |
+| 86 | 2.6893 | 3.94e-05 | 0.45 |
+| 88 | 2.6788 | 3.93e-05 | 0.46 |
+| 90 | 2.6831 | 3.92e-05 | 0.47 |
+| 92 | 2.7039 | 3.91e-05 | 0.48 |
+| 94 | 2.6619 | 3.91e-05 | 0.49 |
+| 96 | 2.6903 | 3.90e-05 | 0.50 |
+| 98 | 2.6993 | 3.89e-05 | 0.51 |
+| 100 | 2.6891 | 3.88e-05 | 0.52 |
+| 102 | 2.6739 | 3.87e-05 | 0.53 |
+
+> **Note:** BitCPM has its own training dataset and data mixture. It is expected that the loss continues to decrease when continue pretraining on open-source datasets.
+
+As shown in the table, the loss gradually decreases from ~2.79 to ~2.67, indicating a stable training process and that the model is learning normally.
+
+## File Description
+
+| File | Description |
+| --- | --- |
+| `train.py` | Training script based on HuggingFace Trainer + DeepSpeed |
+| `run.sh` | Launch script with training hyperparameter configuration |
+| `train_sft.py` | Supervised fine-tuning script based on HuggingFace Trainer + DeepSpeed |
+| `run_sft.sh` | Launch script for SFT with hyperparameter configuration |
+| `ds_config.json` | DeepSpeed ZeRO-3 configuration (with CPU offload) |
+| `ds_config_z2.json` | DeepSpeed ZeRO-2 configuration (used by default) |
+| `requirements.txt` | Python dependency list |

example/ds_config.json

@@ -0,0 +1,29 @@
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "none"
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 2e8,
+        "stage3_prefetch_bucket_size": 2e8,
+        "stage3_param_persistence_threshold": 1e5,
+        "stage3_max_live_parameters": 2e9,
+        "stage3_max_reuse_distance": 2e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

example/ds_config_z2.json

@@ -0,0 +1,22 @@
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "none"
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

example/requirements.txt

@@ -0,0 +1,8 @@
+transformers==4.46.3
+tokenizers==0.20.3
+accelerate==1.1.1
+deepspeed==0.16.2
+datasets==3.1.0
+safetensors==0.4.5
+pyarrow==17.0.0
+tensorboard==2.18.0

example/run.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+MODEL_PATH="/model/BitCPM/BitCPM4-CANN-1B-unquantized/"
+DATA_PATH="/dataset/c4-pro/data/000_1_7.parquet"
+OUTPUT_DIR="./output"
+DS_CONFIG="./ds_config_z2.json"
+
+NUM_GPUS=8
+BATCH_SIZE_PER_GPU=8
+GRAD_ACCUM_STEPS=8
+MAX_SEQ_LENGTH=1024
+
+export ASCEND_RT_VISIBLE_DEVICES=8,9,10,11,12,13,14,15
+
+torchrun --nproc_per_node=$NUM_GPUS train.py \
+    --model_name_or_path $MODEL_PATH \
+    --data_path $DATA_PATH \
+    --max_seq_length $MAX_SEQ_LENGTH \
+    --output_dir $OUTPUT_DIR \
+    --per_device_train_batch_size $BATCH_SIZE_PER_GPU \
+    --gradient_accumulation_steps $GRAD_ACCUM_STEPS \
+    --max_steps 500 \
+    --learning_rate 4e-5 \
+    --lr_scheduler_type cosine \
+    --warmup_ratio 0.1 \
+    --weight_decay 1e-2 \
+    --logging_steps 2 \
+    --save_steps 500 \
+    --save_total_limit 3 \
+    --bf16 \
+    --deepspeed $DS_CONFIG \
+    --gradient_checkpointing \
+    --seed 42 \
+    --dataloader_num_workers 4 \
+    --report_to tensorboard \
+    --logging_dir /data/tensorboard/ \
+    --gradient_checkpointing_kwargs '{"use_reentrant": false}'

example/run_sft.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+
+MODEL_PATH="/model/BitCPM/BitCPM4-CANN-3B-unquantized/"
+DATA_PATH=""
+OUTPUT_DIR="./output_sft"
+DS_CONFIG="./ds_config.json"
+
+NUM_GPUS=8
+BATCH_SIZE_PER_GPU=2
+GRAD_ACCUM_STEPS=1
+MAX_SEQ_LENGTH=4096
+
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+
+torchrun --nproc_per_node=$NUM_GPUS train_sft.py \
+    --model_name_or_path $MODEL_PATH \
+    --data_path $DATA_PATH \
+    --max_seq_length $MAX_SEQ_LENGTH \
+    --output_dir $OUTPUT_DIR \
+    --per_device_train_batch_size $BATCH_SIZE_PER_GPU \
+    --gradient_accumulation_steps $GRAD_ACCUM_STEPS \
+    --num_train_epochs 3 \
+    --learning_rate 2e-5 \
+    --lr_scheduler_type cosine \
+    --warmup_ratio 0.03 \
+    --weight_decay 0.0 \
+    --logging_steps 2 \
+    --save_steps 500 \
+    --save_total_limit 3 \
+    --bf16 \
+    --deepspeed $DS_CONFIG \
+    --gradient_checkpointing \
+    --seed 42 \
+    --dataloader_num_workers 4 \
+    --report_to tensorboard \
+    --logging_dir /data/tensorboard/sft \
+    --train_on_prompt false \
+    --gradient_checkpointing_kwargs '{"use_reentrant": false}'

example/train.py

@@ -0,0 +1,203 @@
+"""
+Continual pretraining script for CPM-2B model using DeepSpeed + HuggingFace Trainer.
+"""
+
+import os
+import json
+import math
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+
+import contextlib
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AutoConfig,
+    Trainer,
+    TrainingArguments,
+    HfArgumentParser,
+    DataCollatorForLanguageModeling,
+    set_seed,
+)
+
+import deepspeed
+_orig_no_sync = deepspeed.DeepSpeedEngine.no_sync
+
+@contextlib.contextmanager
+def _patched_no_sync(self):
+    try:
+        with _orig_no_sync(self):
+            yield
+    except AssertionError:
+        yield
+
+deepspeed.DeepSpeedEngine.no_sync = _patched_no_sync
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier"}
+    )
+    torch_dtype: Optional[str] = field(
+        default="bfloat16",
+        metadata={"help": "torch dtype for model weights (float16, bfloat16, float32)"},
+    )
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(
+        metadata={"help": "Path to training data (parquet file or directory)"}
+    )
+    max_seq_length: int = field(
+        default=4096,
+        metadata={"help": "Maximum sequence length for training"},
+    )
+    text_column: str = field(
+        default="text",
+        metadata={"help": "Name of the text column in the dataset"},
+    )
+    preprocessing_num_workers: int = field(
+        default=8,
+        metadata={"help": "Number of workers for data preprocessing"},
+    )
+
+
+def tokenize_and_group(dataset, tokenizer, data_args):
+    """Tokenize texts and group into chunks of max_seq_length."""
+
+    column_names = dataset.column_names
+    text_column = data_args.text_column
+    if text_column not in column_names:
+        candidates = [c for c in column_names if "text" in c.lower()]
+        if candidates:
+            text_column = candidates[0]
+        else:
+            text_column = column_names[0]
+        logger.warning(f"Column '{data_args.text_column}' not found, using '{text_column}'")
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column], add_special_tokens=False)
+
+    tokenized_dataset = dataset.map(
+        tokenize_function,
+        batched=True,
+        num_proc=data_args.preprocessing_num_workers,
+        remove_columns=column_names,
+        desc="Tokenizing",
+    )
+
+    block_size = data_args.max_seq_length
+
+    def group_texts(examples):
+        concatenated = {k: sum(examples[k], []) for k in examples.keys()}
+        total_length = len(concatenated["input_ids"])
+        total_length = (total_length // block_size) * block_size
+
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    grouped_dataset = tokenized_dataset.map(
+        group_texts,
+        batched=True,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="Grouping texts",
+    )
+
+    return grouped_dataset
+
+
+def main():
+    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
+    )
+    logger.info(f"Training args: {training_args}")
+
+    set_seed(training_args.seed)
+
+    dtype_map = {
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+    }
+    torch_dtype = dtype_map.get(model_args.torch_dtype, torch.bfloat16)
+
+    logger.info(f"Loading tokenizer from {model_args.model_name_or_path}")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        trust_remote_code=True,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    logger.info(f"Loading model from {model_args.model_name_or_path}")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        torch_dtype=torch_dtype,
+        trust_remote_code=True,
+        attn_implementation="sdpa",
+    )
+    model.config.use_cache = False
+
+    logger.info(f"Loading dataset from {data_args.data_path}")
+    if os.path.isfile(data_args.data_path):
+        raw_dataset = load_dataset("parquet", data_files=data_args.data_path, split="train")
+    elif os.path.isdir(data_args.data_path):
+        parquet_files = [
+            os.path.join(data_args.data_path, f)
+            for f in os.listdir(data_args.data_path)
+            if f.endswith(".parquet")
+        ]
+        raw_dataset = load_dataset("parquet", data_files=parquet_files, split="train")
+    else:
+        raise ValueError(f"Data path not found: {data_args.data_path}")
+
+    logger.info(f"Dataset loaded: {len(raw_dataset)} samples, columns: {raw_dataset.column_names}")
+
+    train_dataset = tokenize_and_group(raw_dataset, tokenizer, data_args)
+    logger.info(f"Processed dataset: {len(train_dataset)} samples of length {data_args.max_seq_length}")
+
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=data_collator,
+    )
+
+    logger.info("Starting training...")
+    train_result = trainer.train(
+        resume_from_checkpoint=training_args.resume_from_checkpoint
+    )
+
+    trainer.save_model()
+    trainer.save_state()
+
+    metrics = train_result.metrics
+    metrics["train_samples"] = len(train_dataset)
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+
+
+if __name__ == "__main__":
+    main()

example/train_sft.py

@@ -0,0 +1,424 @@
+"""
+Supervised fine-tuning script using DeepSpeed + HuggingFace Trainer.
+"""
+
+import json
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import contextlib
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+
+import deepspeed
+_orig_no_sync = deepspeed.DeepSpeedEngine.no_sync
+
+@contextlib.contextmanager
+def _patched_no_sync(self):
+    try:
+        with _orig_no_sync(self):
+            yield
+    except AssertionError:
+        yield
+
+deepspeed.DeepSpeedEngine.no_sync = _patched_no_sync
+
+logger = logging.getLogger(__name__)
+
+IGNORE_INDEX = -100
+
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier"}
+    )
+    torch_dtype: Optional[str] = field(
+        default="bfloat16",
+        metadata={"help": "torch dtype for model weights (float16, bfloat16, float32)"},
+    )
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(metadata={"help": "Path to SFT data file or directory"})
+    max_seq_length: int = field(
+        default=4096,
+        metadata={"help": "Maximum sequence length for training"},
+    )
+    prompt_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "Prompt/instruction column name. Auto-detected if omitted."},
+    )
+    input_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "Optional extra input/context column name"},
+    )
+    response_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "Response/output column name. Auto-detected if omitted."},
+    )
+    messages_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "Chat messages column name. Auto-detected if omitted."},
+    )
+    system_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "Optional system prompt column name"},
+    )
+    train_on_prompt: bool = field(
+        default=False,
+        metadata={"help": "Whether to compute loss on prompt/user tokens"},
+    )
+    add_eos_token: bool = field(
+        default=True,
+        metadata={"help": "Append eos_token to plain prompt/response examples"},
+    )
+    preprocessing_num_workers: int = field(
+        default=8,
+        metadata={"help": "Number of workers for data preprocessing"},
+    )
+
+
+class SFTDataCollator:
+    def __init__(self, tokenizer, pad_to_multiple_of: Optional[int] = 8):
+        self.tokenizer = tokenizer
+        self.pad_to_multiple_of = pad_to_multiple_of
+
+    def __call__(self, features: List[Dict[str, List[int]]]) -> Dict[str, torch.Tensor]:
+        max_length = max(len(feature["input_ids"]) for feature in features)
+        if self.pad_to_multiple_of:
+            multiple = self.pad_to_multiple_of
+            max_length = ((max_length + multiple - 1) // multiple) * multiple
+
+        input_ids = []
+        attention_mask = []
+        labels = []
+        pad_token_id = self.tokenizer.pad_token_id
+
+        for feature in features:
+            length = len(feature["input_ids"])
+            pad_length = max_length - length
+            input_ids.append(feature["input_ids"] + [pad_token_id] * pad_length)
+            attention_mask.append([1] * length + [0] * pad_length)
+            labels.append(feature["labels"] + [IGNORE_INDEX] * pad_length)
+
+        return {
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+        }
+
+
+def load_sft_dataset(data_path: str):
+    if os.path.isfile(data_path):
+        extension = os.path.splitext(data_path)[1].lstrip(".").lower()
+        if extension == "jsonl":
+            extension = "json"
+        if extension not in {"parquet", "json", "csv", "txt"}:
+            raise ValueError(f"Unsupported data file extension: {extension}")
+        return load_dataset(extension, data_files=data_path, split="train")
+
+    if os.path.isdir(data_path):
+        data_files = []
+        extension = None
+        for name in os.listdir(data_path):
+            current_extension = os.path.splitext(name)[1].lstrip(".").lower()
+            if current_extension == "jsonl":
+                current_extension = "json"
+            if current_extension in {"parquet", "json", "csv", "txt"}:
+                extension = extension or current_extension
+                if current_extension == extension:
+                    data_files.append(os.path.join(data_path, name))
+        if not data_files or extension is None:
+            raise ValueError(f"No supported data files found in: {data_path}")
+        return load_dataset(extension, data_files=sorted(data_files), split="train")
+
+    raise ValueError(f"Data path not found: {data_path}")
+
+
+def choose_column(
+    column_names: List[str], explicit: Optional[str], candidates: List[str]
+) -> Optional[str]:
+    if explicit:
+        if explicit not in column_names:
+            raise ValueError(f"Column '{explicit}' not found. Available columns: {column_names}")
+        return explicit
+    for name in candidates:
+        if name in column_names:
+            return name
+    return None
+
+
+def parse_messages(value: Any) -> List[Dict[str, str]]:
+    if isinstance(value, str):
+        value = json.loads(value)
+    if not isinstance(value, list):
+        raise ValueError("messages/conversations column must be a list or JSON string")
+
+    messages = []
+    for item in value:
+        if not isinstance(item, dict):
+            raise ValueError("Each message must be a dict")
+
+        role = item.get("role", item.get("from"))
+        content = item.get("content", item.get("value"))
+        if role == "human":
+            role = "user"
+        elif role == "gpt":
+            role = "assistant"
+
+        if role is None or content is None:
+            raise ValueError("Each message must contain role/from and content/value")
+        messages.append({"role": str(role), "content": str(content)})
+
+    return messages
+
+
+def tokenize_text(tokenizer, text: str) -> List[int]:
+    return tokenizer(text, add_special_tokens=False)["input_ids"]
+
+
+def apply_chat_template(tokenizer, messages: List[Dict[str, str]], add_generation_prompt: bool) -> str:
+    if tokenizer.chat_template is None:
+        raise ValueError(
+            "The tokenizer has no chat_template. Use prompt/response columns or set a chat_template."
+        )
+    return tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=add_generation_prompt,
+    )
+
+
+def encode_prompt_response(
+    example: Dict[str, Any],
+    tokenizer,
+    data_args: DataArguments,
+    prompt_column: str,
+    input_column: Optional[str],
+    response_column: str,
+) -> Tuple[List[int], List[int]]:
+    prompt = str(example[prompt_column])
+    if input_column and example.get(input_column):
+        prompt = prompt + "\n" + str(example[input_column])
+    response = str(example[response_column])
+
+    messages = []
+    if data_args.system_column and example.get(data_args.system_column):
+        messages.append({"role": "system", "content": str(example[data_args.system_column])})
+    messages.append({"role": "user", "content": prompt})
+    messages.append({"role": "assistant", "content": response})
+
+    if tokenizer.chat_template is not None:
+        full_text = apply_chat_template(tokenizer, messages, add_generation_prompt=False)
+        prompt_text = apply_chat_template(tokenizer, messages[:-1], add_generation_prompt=True)
+        input_ids = tokenize_text(tokenizer, full_text)
+        prompt_length = len(tokenize_text(tokenizer, prompt_text))
+    else:
+        response_text = response
+        if data_args.add_eos_token and tokenizer.eos_token:
+            response_text += tokenizer.eos_token
+        full_text = prompt + "\n" + response_text
+        input_ids = tokenize_text(tokenizer, full_text)
+        prompt_length = len(tokenize_text(tokenizer, prompt + "\n"))
+
+    labels = input_ids.copy()
+    if not data_args.train_on_prompt:
+        labels[:prompt_length] = [IGNORE_INDEX] * min(prompt_length, len(labels))
+    return input_ids, labels
+
+
+def encode_messages(
+    example: Dict[str, Any],
+    tokenizer,
+    data_args: DataArguments,
+    messages_column: str,
+) -> Tuple[List[int], List[int]]:
+    messages = parse_messages(example[messages_column])
+
+    if tokenizer.chat_template is not None:
+        full_text = apply_chat_template(tokenizer, messages, add_generation_prompt=False)
+        input_ids = tokenize_text(tokenizer, full_text)
+        labels = [IGNORE_INDEX] * len(input_ids)
+
+        if data_args.train_on_prompt:
+            labels = input_ids.copy()
+        else:
+            for index, message in enumerate(messages):
+                if message["role"] != "assistant":
+                    continue
+                before_text = apply_chat_template(
+                    tokenizer, messages[:index], add_generation_prompt=True
+                )
+                after_text = apply_chat_template(
+                    tokenizer, messages[: index + 1], add_generation_prompt=False
+                )
+                start = len(tokenize_text(tokenizer, before_text))
+                end = len(tokenize_text(tokenizer, after_text))
+                labels[start:end] = input_ids[start:end]
+    else:
+        labels = []
+        input_ids = []
+        for message in messages:
+            part = f"{message['role']}: {message['content']}\n"
+            if data_args.add_eos_token and message["role"] == "assistant" and tokenizer.eos_token:
+                part += tokenizer.eos_token
+            part_ids = tokenize_text(tokenizer, part)
+            input_ids.extend(part_ids)
+            if data_args.train_on_prompt or message["role"] == "assistant":
+                labels.extend(part_ids)
+            else:
+                labels.extend([IGNORE_INDEX] * len(part_ids))
+
+    return input_ids, labels
+
+
+def preprocess_sft_dataset(raw_dataset, tokenizer, data_args: DataArguments):
+    column_names = raw_dataset.column_names
+    messages_column = choose_column(
+        column_names, data_args.messages_column, ["messages", "conversations"]
+    )
+    prompt_column = choose_column(
+        column_names,
+        data_args.prompt_column,
+        ["prompt", "instruction", "question"],
+    )
+    input_column = choose_column(
+        column_names,
+        data_args.input_column,
+        ["input", "context"],
+    )
+    response_column = choose_column(
+        column_names,
+        data_args.response_column,
+        ["response", "output", "answer", "chosen"],
+    )
+
+    if messages_column:
+        logger.info(f"Using chat messages column: {messages_column}")
+    elif prompt_column and response_column:
+        logger.info(f"Using prompt column '{prompt_column}' and response column '{response_column}'")
+    else:
+        raise ValueError(
+            "Cannot infer SFT data format. Provide either messages/conversations or "
+            "prompt/instruction plus response/output columns."
+        )
+
+    def encode_batch(examples):
+        batch_input_ids = []
+        batch_labels = []
+        batch_attention_mask = []
+
+        batch_size = len(next(iter(examples.values())))
+        for i in range(batch_size):
+            example = {name: values[i] for name, values in examples.items()}
+            if messages_column:
+                input_ids, labels = encode_messages(example, tokenizer, data_args, messages_column)
+            else:
+                input_ids, labels = encode_prompt_response(
+                    example, tokenizer, data_args, prompt_column, input_column, response_column
+                )
+
+            input_ids = input_ids[: data_args.max_seq_length]
+            labels = labels[: data_args.max_seq_length]
+            if not input_ids or all(label == IGNORE_INDEX for label in labels):
+                continue
+
+            batch_input_ids.append(input_ids)
+            batch_labels.append(labels)
+            batch_attention_mask.append([1] * len(input_ids))
+
+        return {
+            "input_ids": batch_input_ids,
+            "attention_mask": batch_attention_mask,
+            "labels": batch_labels,
+        }
+
+    return raw_dataset.map(
+        encode_batch,
+        batched=True,
+        num_proc=data_args.preprocessing_num_workers,
+        remove_columns=column_names,
+        desc="Tokenizing SFT data",
+    )
+
+
+def main():
+    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
+    )
+    logger.info(f"Training args: {training_args}")
+
+    set_seed(training_args.seed)
+
+    dtype_map = {
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+    }
+    torch_dtype = dtype_map.get(model_args.torch_dtype, torch.bfloat16)
+
+    logger.info(f"Loading tokenizer from {model_args.model_name_or_path}")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        trust_remote_code=True,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    logger.info(f"Loading model from {model_args.model_name_or_path}")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        torch_dtype=torch_dtype,
+        trust_remote_code=True,
+        attn_implementation="sdpa",
+    )
+    model.config.use_cache = False
+
+    logger.info(f"Loading SFT dataset from {data_args.data_path}")
+    raw_dataset = load_sft_dataset(data_args.data_path)
+    logger.info(f"Dataset loaded: {len(raw_dataset)} samples, columns: {raw_dataset.column_names}")
+
+    train_dataset = preprocess_sft_dataset(raw_dataset, tokenizer, data_args)
+    logger.info(f"Processed dataset: {len(train_dataset)} samples")
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=SFTDataCollator(tokenizer),
+    )
+
+    logger.info("Starting SFT training...")
+    train_result = trainer.train(
+        resume_from_checkpoint=training_args.resume_from_checkpoint
+    )
+
+    trainer.save_model()
+    trainer.save_state()
+
+    metrics = train_result.metrics
+    metrics["train_samples"] = len(train_dataset)
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+
+
+if __name__ == "__main__":
+    main()