Remove misplaced files from root (moved to example/)

ds_config.json

@@ -1,29 +0,0 @@
-{
-    "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
-}

ds_config_z2.json

@@ -1,22 +0,0 @@
-{
-    "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
-}

requirements.txt

@@ -1,8 +0,0 @@
-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

run.sh

@@ -1,37 +0,0 @@
-#!/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}'

run_sft.sh

@@ -1,38 +0,0 @@
-#!/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}'

train.py

@@ -1,203 +0,0 @@
-"""
-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()

train_sft.py

@@ -1,424 +0,0 @@
-"""
-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()