Upload scripts/train_n8n_sft.py with huggingface_hub

scripts/train_n8n_sft.py

@@ -0,0 +1,275 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "transformers>=4.45.0",
+#     "trl>=0.12.0",
+#     "peft>=0.13.0",
+#     "datasets>=3.0.0",
+#     "accelerate>=1.0.0",
+#     "bitsandbytes>=0.44.0",
+#     "wandb>=0.18.0",
+#     "huggingface_hub>=0.26.0",
+#     "torch>=2.4.0",
+#     "einops>=0.8.0",
+#     "sentencepiece>=0.2.0",
+# ]
+# [tool.uv]
+# extra-index-url = ["https://download.pytorch.org/whl/cu124"]
+# ///
+"""
+Script d'entraînement SFT pour le modèle n8n Expert.
+
+Usage sur HuggingFace Jobs:
+    hf jobs uv run \
+        --script train_n8n_sft.py \
+        --flavor h100x1 \
+        --name n8n-expert-sft \
+        --timeout 24h
+
+Variables d'environnement requises:
+    - HF_TOKEN: Token HuggingFace avec accès en écriture
+    - WANDB_API_KEY: (optionnel) Pour le tracking W&B
+"""
+
+import os
+import json
+import torch
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from trl import SFTTrainer, SFTConfig
+from huggingface_hub import login
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+# Modèle de base
+MODEL_NAME = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-14B-Instruct")
+
+# Dataset
+DATASET_REPO = "stmasson/n8n-agentic-multitask"
+TRAIN_FILE = "data/multitask_large/train.jsonl"
+VAL_FILE = "data/multitask_large/val.jsonl"
+
+# Output
+OUTPUT_DIR = "./n8n-expert-sft"
+HF_REPO = os.environ.get("HF_REPO", "stmasson/n8n-expert-14b-sft")
+
+# Hyperparamètres
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", "3"))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "2"))
+GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "8"))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "2e-5"))
+MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", "8192"))
+
+# LoRA
+LORA_R = int(os.environ.get("LORA_R", "64"))
+LORA_ALPHA = int(os.environ.get("LORA_ALPHA", "128"))
+LORA_DROPOUT = float(os.environ.get("LORA_DROPOUT", "0.05"))
+
+# Quantization (pour économiser la VRAM)
+USE_4BIT = os.environ.get("USE_4BIT", "false").lower() == "true"
+
+# ============================================================================
+# AUTHENTIFICATION
+# ============================================================================
+
+print("=" * 60)
+print("ENTRAÎNEMENT SFT - N8N EXPERT")
+print("=" * 60)
+
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
+    print("Authentifié sur HuggingFace")
+else:
+    print("Warning: HF_TOKEN non défini, push désactivé")
+
+wandb_key = os.environ.get("WANDB_API_KEY")
+if wandb_key:
+    import wandb
+    wandb.login(key=wandb_key)
+    report_to = "wandb"
+    print("Tracking W&B activé")
+else:
+    report_to = "none"
+    print("Tracking W&B désactivé")
+
+# ============================================================================
+# CHARGEMENT DU MODÈLE
+# ============================================================================
+
+print(f"\nChargement du modèle: {MODEL_NAME}")
+
+# Configuration quantization si nécessaire
+if USE_4BIT:
+    print("Mode 4-bit activé")
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    model = prepare_model_for_kbit_training(model)
+else:
+    print("Mode bfloat16")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+# Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+print(f"Modèle chargé: {model.config.num_hidden_layers} layers, {model.config.hidden_size} hidden size")
+
+# ============================================================================
+# CONFIGURATION LORA
+# ============================================================================
+
+print(f"\nConfiguration LoRA: r={LORA_R}, alpha={LORA_ALPHA}")
+
+lora_config = LoraConfig(
+    r=LORA_R,
+    lora_alpha=LORA_ALPHA,
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj"
+    ],
+    lora_dropout=LORA_DROPOUT,
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+
+# ============================================================================
+# CHARGEMENT DU DATASET
+# ============================================================================
+
+print(f"\nChargement du dataset: {DATASET_REPO}")
+
+dataset = load_dataset(
+    DATASET_REPO,
+    data_files={
+        "train": TRAIN_FILE,
+        "validation": VAL_FILE
+    }
+)
+
+print(f"Train: {len(dataset['train'])} exemples")
+print(f"Validation: {len(dataset['validation'])} exemples")
+
+# Fonction de formatage
+def format_example(example):
+    """Formate les messages en texte pour l'entraînement"""
+    messages = example["messages"]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=False
+    )
+    return {"text": text}
+
+# Appliquer le formatage
+print("Formatage des données...")
+dataset = dataset.map(format_example, remove_columns=dataset["train"].column_names)
+
+# Afficher un exemple
+print("\nExemple de données formatées:")
+print(dataset["train"][0]["text"][:500] + "...")
+
+# ============================================================================
+# CONFIGURATION D'ENTRAÎNEMENT
+# ============================================================================
+
+print(f"\nConfiguration d'entraînement:")
+print(f"  - Epochs: {NUM_EPOCHS}")
+print(f"  - Batch size: {BATCH_SIZE}")
+print(f"  - Gradient accumulation: {GRAD_ACCUM}")
+print(f"  - Effective batch size: {BATCH_SIZE * GRAD_ACCUM}")
+print(f"  - Learning rate: {LEARNING_RATE}")
+print(f"  - Max sequence length: {MAX_SEQ_LENGTH}")
+
+training_args = SFTConfig(
+    output_dir=OUTPUT_DIR,
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=BATCH_SIZE,
+    per_device_eval_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LEARNING_RATE,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.1,
+    weight_decay=0.01,
+    bf16=True,
+    tf32=True,
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=500,
+    save_total_limit=3,
+    eval_strategy="steps",
+    eval_steps=500,
+    max_seq_length=MAX_SEQ_LENGTH,
+    packing=True,
+    gradient_checkpointing=True,
+    gradient_checkpointing_kwargs={"use_reentrant": False},
+    dataset_text_field="text",
+    report_to=report_to,
+    run_name="n8n-expert-sft",
+    hub_model_id=HF_REPO if hf_token else None,
+    push_to_hub=bool(hf_token),
+    hub_strategy="checkpoint",
+)
+
+# ============================================================================
+# ENTRAÎNEMENT
+# ============================================================================
+
+print("\nInitialisation du trainer...")
+
+trainer = SFTTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["validation"],
+    peft_config=lora_config,
+    tokenizer=tokenizer,
+)
+
+# Afficher les paramètres entraînables
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total_params = sum(p.numel() for p in model.parameters())
+print(f"\nParamètres entraînables: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")
+
+print("\n" + "=" * 60)
+print("DÉMARRAGE DE L'ENTRAÎNEMENT")
+print("=" * 60)
+
+trainer.train()
+
+# ============================================================================
+# SAUVEGARDE
+# ============================================================================
+
+print("\nSauvegarde du modèle...")
+trainer.save_model(f"{OUTPUT_DIR}/final")
+
+if hf_token:
+    print(f"Push vers {HF_REPO}...")
+    trainer.push_to_hub()
+    print(f"Modèle disponible sur: https://huggingface.co/{HF_REPO}")
+
+print("\n" + "=" * 60)
+print("ENTRAÎNEMENT TERMINÉ")
+print("=" * 60)