大模型实践--多显卡分布式训练实践（基于DeepSpeed）

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  gradient_accumulation_steps: 8
  gradient_clipping: 1.0
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: false
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
import torch
from transformers import AutoModelForCausalLM,AutoTokenizer
import os

# 配置预训练模型下载地址
os.environ['HF_ENDPOINT']='https://hf-mirror.com/'
os.environ['HF_HOME']='/root/autodl-tmp/hf'

# Configure model and tokenizer
model_name = "Qwen/Qwen3-4B"
model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name,dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)

# 数据集
# Load dataset
dataset_dict = load_dataset("json",data_files="/root/autodl-tmp/hf/data/psychology_data.jsonl")
dataset_dict = dataset_dict['train'].train_test_split(test_size=0.1,shuffle=True)

def map_func(message):
    conversation = message['conversation']
    messages = []
    for item in conversation:
        messages.append({"role": "user", "content": item['human']})
        messages.append({"role": "assistant", "content":item['assistant']})
    return {"messages":messages}
dataset_dict = dataset_dict.map(map_func,batched=False,remove_columns=['conversation_id','category','conversation','dataset'])

# Configure trainer
training_args = SFTConfig(
    output_dir="/root/autodl-tmp/hf/model/Qwen3-4B/full_zero3/model1/",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=8,
    learning_rate=5e-5,
    save_steps=20,
    save_total_limit=1,
    eval_strategy="steps",
    eval_steps=20,
    load_best_model_at_end=True,
    logging_dir="/root/tf-logs/",
    logging_strategy='steps',
    logging_steps=10,
    bf16=True,
    warmup_ratio=0.1,
)
# Initialize trainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset_dict["train"],
    eval_dataset=dataset_dict["test"],
    processing_class=tokenizer,
)

# Start training
trainer.train()

from transformers import AutoModelForCausalLM, AutoTokenizer,TextStreamer

model_name = "Qwen/Qwen3-4B"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype="auto",
    device_map="auto"
)

prompt = "觉得找心里咨询师很丢脸? 我今年17，因为某些原因觉得要找个心理咨询师咨询一下，但我没法向家长开口，我觉得心理问题就像得了艾滋病一样让人难以启齿，很难堪。而且父母也不会理解，现在花一样的年纪，能有什么心坎。我该怎么办?"
inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
streamer = TextStreamer(tokenizer,skip_prompt=True)

# Despite returning the usual output, the streamer will also print the generated text to stdout.
_ = model.generate(**inputs, streamer=streamer, max_new_tokens=1024)

from transformers import AutoModelForCausalLM, AutoTokenizer,TextStreamer

model_name = "/root/autodl-tmp/hf/model/Qwen3-4B/full_zero3/model1/checkpoint-272/"
# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype="auto",
    device_map="auto"
)

# # prepare the model input
prompt = "觉得找心里咨询师很丢脸? 我今年17，因为某些原因觉得要找个心理咨询师咨询一下，但我没法向家长开口，我觉得心理问题就像得了艾滋病一样让人难以启齿，很难堪。而且父母也不会理解，现在花一样的年纪，能有什么心坎。我该怎么办?"
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False # Switches between thinking and non-thinking modes. Default is True.
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# conduct text completion
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=32768
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# parsing thinking content
try:
    # rindex finding 151668 (</think>)
    index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:
    index = 0

thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

print("thinking content:", thinking_content)
print("content:", content)