Fix: use TrainingArguments instead of SFTConfig

train.py

@@ -0,0 +1,397 @@
+"""
+D1337 CIPHER - Custom Training Script
+=====================================
+Optimized QLoRA training for 31B model on 4x L40S (192GB VRAM)
+
+Brand: D1337 SOVEREIGN LABS
+Model: GLM-4.7-Flash-abliterated (31B) -> D1337 CIPHER
+"""
+
+import os
+import sys
+import torch
+import gradio as gr
+from threading import Thread
+from dataclasses import dataclass
+from typing import Optional
+
+# Training imports
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    BitsAndBytesConfig,
+)
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    TaskType,
+)
+from datasets import load_dataset
+from trl import SFTTrainer, SFTConfig
+
+# ============================================
+# CONFIGURATION
+# ============================================
+@dataclass
+class TrainingConfig:
+    # Model
+    base_model: str = "huihui-ai/Huihui-GLM-4.7-Flash-abliterated"
+    output_model: str = "Desorden1337/d1337-cipher-v1"
+    
+    # Dataset
+    dataset_name: str = "Desorden1337/d1337-cipher-dataset"
+    dataset_split: str = "train"
+    
+    # LoRA Config (reduced for 4x L40S memory)
+    lora_r: int = 32
+    lora_alpha: int = 64
+    lora_dropout: float = 0.05
+    target_modules: list = None
+    
+    # Training
+    num_epochs: int = 5
+    batch_size: int = 1
+    gradient_accumulation: int = 8
+    learning_rate: float = 2e-4
+    max_seq_length: int = 2048  # Reduced for memory
+    warmup_ratio: float = 0.1
+    weight_decay: float = 0.01
+    
+    # Hardware
+    use_4bit: bool = True
+    use_bf16: bool = True
+    
+    def __post_init__(self):
+        if self.target_modules is None:
+            self.target_modules = [
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj"
+            ]
+
+
+# ============================================
+# TRAINING CLASS
+# ============================================
+class D1337CipherTrainer:
+    def __init__(self, config: TrainingConfig = None):
+        self.config = config or TrainingConfig()
+        self.model = None
+        self.tokenizer = None
+        self.trainer = None
+        self.training_status = "Idle"
+        self.training_log = []
+        
+    def log(self, message: str):
+        """Log message to console and internal log"""
+        print(f"[D1337] {message}")
+        self.training_log.append(message)
+        if len(self.training_log) > 100:
+            self.training_log = self.training_log[-100:]
+    
+    def setup_quantization(self):
+        """Setup 4-bit quantization config"""
+        if self.config.use_4bit:
+            return BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16 if self.config.use_bf16 else torch.float16,
+                bnb_4bit_use_double_quant=True,
+            )
+        return None
+    
+    def setup_lora(self):
+        """Setup LoRA configuration"""
+        return LoraConfig(
+            r=self.config.lora_r,
+            lora_alpha=self.config.lora_alpha,
+            lora_dropout=self.config.lora_dropout,
+            target_modules=self.config.target_modules,
+            bias="none",
+            task_type=TaskType.CAUSAL_LM,
+        )
+    
+    def load_model(self):
+        """Load base model with quantization"""
+        self.training_status = "Loading model..."
+        self.log(f"Loading model: {self.config.base_model}")
+        
+        # Load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.base_model,
+            trust_remote_code=True,
+            padding_side="right",
+        )
+        
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        # Load model with quantization
+        bnb_config = self.setup_quantization()
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.config.base_model,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16 if self.config.use_bf16 else torch.float16,
+        )
+        
+        # Enable gradient checkpointing for memory efficiency
+        self.model.gradient_checkpointing_enable()
+        self.model.enable_input_require_grads()
+        
+        # Apply LoRA
+        lora_config = self.setup_lora()
+        self.model = get_peft_model(self.model, lora_config)
+        
+        # Print trainable parameters
+        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        total_params = sum(p.numel() for p in self.model.parameters())
+        self.log(f"Trainable parameters: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")
+        self.log(f"Model loaded on {torch.cuda.device_count()} GPU(s)")
+        
+    def load_dataset(self):
+        """Load and prepare dataset"""
+        self.training_status = "Loading dataset..."
+        self.log(f"Loading dataset: {self.config.dataset_name}")
+        
+        dataset = load_dataset(self.config.dataset_name, split=self.config.dataset_split)
+        self.log(f"Dataset loaded: {len(dataset)} samples")
+        
+        return dataset
+    
+    def format_messages(self, example):
+        """Format messages into training text"""
+        messages = example["messages"]
+        
+        # Use ChatML format
+        text = ""
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            text += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+        
+        return {"text": text}
+    
+    def train(self):
+        """Execute training"""
+        try:
+            self.training_status = "Initializing..."
+            self.log("=" * 60)
+            self.log("D1337 CIPHER TRAINING - INITIATED")
+            self.log("=" * 60)
+            
+            # Load model and dataset
+            self.load_model()
+            dataset = self.load_dataset()
+            
+            # Format dataset
+            self.log("Formatting dataset...")
+            dataset = dataset.map(self.format_messages, remove_columns=dataset.column_names)
+            
+            # Training arguments (standard TrainingArguments)
+            self.training_status = "Setting up training..."
+            training_args = TrainingArguments(
+                output_dir="./d1337-cipher-output",
+                num_train_epochs=self.config.num_epochs,
+                per_device_train_batch_size=self.config.batch_size,
+                gradient_accumulation_steps=self.config.gradient_accumulation,
+                learning_rate=self.config.learning_rate,
+                weight_decay=self.config.weight_decay,
+                warmup_steps=14,
+                lr_scheduler_type="cosine",
+                logging_steps=1,
+                save_steps=50,
+                save_total_limit=2,
+                bf16=self.config.use_bf16,
+                fp16=not self.config.use_bf16,
+                gradient_checkpointing=True,
+                max_grad_norm=1.0,
+                group_by_length=True,
+                dataloader_num_workers=4,
+                remove_unused_columns=False,
+                push_to_hub=True,
+                hub_model_id=self.config.output_model,
+                hub_private_repo=True,
+                report_to="none",
+            )
+            
+            # Initialize trainer with explicit tokenizer
+            self.trainer = SFTTrainer(
+                model=self.model,
+                args=training_args,
+                train_dataset=dataset,
+            )
+            
+            # Start training
+            self.training_status = "Training in progress..."
+            self.log("Training started!")
+            self.trainer.train()
+            
+            # Save and push
+            self.training_status = "Saving model..."
+            self.log("Saving model...")
+            self.trainer.save_model()
+            self.trainer.push_to_hub()
+            
+            self.training_status = "Complete!"
+            self.log("=" * 60)
+            self.log("D1337 CIPHER TRAINING - COMPLETE!")
+            self.log(f"Model saved to: {self.config.output_model}")
+            self.log("=" * 60)
+            
+            return True
+            
+        except Exception as e:
+            self.training_status = f"Error: {str(e)}"
+            self.log(f"Training failed: {str(e)}")
+            import traceback
+            self.log(traceback.format_exc())
+            return False
+
+
+# ============================================
+# GRADIO UI
+# ============================================
+def create_ui(trainer: D1337CipherTrainer):
+    """Create Gradio UI for monitoring"""
+    
+    def get_status():
+        return trainer.training_status
+    
+    def get_logs():
+        return "\n".join(trainer.training_log[-50:])
+    
+    def start_training():
+        trainer.training_log = []
+        thread = Thread(target=trainer.train)
+        thread.start()
+        return "Training started! Check logs for progress."
+    
+    def get_gpu_info():
+        if torch.cuda.is_available():
+            info = []
+            for i in range(torch.cuda.device_count()):
+                props = torch.cuda.get_device_properties(i)
+                mem_total = props.total_memory / (1024**3)
+                mem_used = torch.cuda.memory_allocated(i) / (1024**3)
+                info.append(f"GPU {i}: {props.name} - {mem_used:.1f}GB / {mem_total:.1f}GB")
+            return "\n".join(info)
+        return "No GPU available"
+    
+    with gr.Blocks(title="D1337 CIPHER Training", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🔥 D1337 CIPHER - Training Console
+        ### D1337 SOVEREIGN LABS
+        
+        Custom training environment for GLM-4.7-Flash-abliterated → D1337 CIPHER
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### Configuration")
+                model_name = gr.Textbox(
+                    label="Base Model",
+                    value=trainer.config.base_model,
+                    interactive=False
+                )
+                dataset_name = gr.Textbox(
+                    label="Dataset",
+                    value=trainer.config.dataset_name,
+                    interactive=False
+                )
+                output_name = gr.Textbox(
+                    label="Output Model",
+                    value=trainer.config.output_model,
+                    interactive=False
+                )
+                
+                gr.Markdown("### Training Parameters")
+                gr.Textbox(
+                    label="LoRA Rank",
+                    value=str(trainer.config.lora_r),
+                    interactive=False
+                )
+                gr.Textbox(
+                    label="Epochs",
+                    value=str(trainer.config.num_epochs),
+                    interactive=False
+                )
+                gr.Textbox(
+                    label="Learning Rate",
+                    value=str(trainer.config.learning_rate),
+                    interactive=False
+                )
+                
+            with gr.Column(scale=2):
+                gr.Markdown("### Status")
+                status_box = gr.Textbox(
+                    label="Current Status",
+                    value=get_status,
+                    every=2
+                )
+                
+                gpu_info = gr.Textbox(
+                    label="GPU Info",
+                    value=get_gpu_info,
+                    every=5
+                )
+                
+                start_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+                
+                gr.Markdown("### Training Logs")
+                logs_box = gr.Textbox(
+                    label="Logs",
+                    value=get_logs,
+                    every=3,
+                    lines=15,
+                    max_lines=20
+                )
+        
+        start_btn.click(fn=start_training, outputs=status_box)
+    
+    return demo
+
+
+# ============================================
+# MAIN
+# ============================================
+def main():
+    print("=" * 60)
+    print("D1337 CIPHER - Custom Training Environment")
+    print("D1337 SOVEREIGN LABS")
+    print("=" * 60)
+    
+    # Check GPU
+    if torch.cuda.is_available():
+        print(f"GPUs available: {torch.cuda.device_count()}")
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            print(f"  GPU {i}: {props.name} ({props.total_memory / (1024**3):.1f} GB)")
+    else:
+        print("WARNING: No GPU detected!")
+    
+    # Initialize trainer
+    config = TrainingConfig()
+    trainer = D1337CipherTrainer(config)
+    
+    # Check if auto-start
+    auto_start = os.environ.get("AUTO_START_TRAINING", "false").lower() == "true"
+    
+    if auto_start:
+        print("Auto-starting training...")
+        trainer.train()
+    else:
+        # Launch Gradio UI
+        print("Launching Gradio UI...")
+        demo = create_ui(trainer)
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False
+        )
+
+
+if __name__ == "__main__":
+    main()