Upload folder using huggingface_hub

run_transformers_training.py

@@ -184,227 +184,291 @@ def load_configs(base_path):
         raise
 
 def parse_args():
-    parser = argparse.ArgumentParser(description="Fine-tune a language model on a text dataset")
-    parser.add_argument("--config", type=str, default="transformers_config.json", help="Path to configuration file")
+    """
+    Parse command line arguments for the training script.
+    
+    Returns:
+        argparse.Namespace: The parsed command line arguments
+    """
+    parser = argparse.ArgumentParser(description="Run training for language models")
+    parser.add_argument(
+        "--config_file", 
+        type=str, 
+        default=None,
+        help="Path to the configuration file (default: transformers_config.json in script directory)"
+    )
+    parser.add_argument(
+        "--seed", 
+        type=int, 
+        default=None,
+        help="Random seed for reproducibility (default: based on current time)"
+    )
+    parser.add_argument(
+        "--log_level", 
+        type=str, 
+        choices=["debug", "info", "warning", "error", "critical"],
+        default="info",
+        help="Logging level (default: info)"
+    )
     return parser.parse_args()
 
 def load_model_and_tokenizer(config):
-    """Load model and tokenizer with proper error handling and optimizations."""
+    """
+    Load the model and tokenizer according to the configuration.
+    
+    Args:
+        config (dict): Complete configuration dictionary
+    
+    Returns:
+        tuple: (model, tokenizer) - The loaded model and tokenizer
+    """
+    # Extract model configuration
+    model_config = get_config_value(config, "model", {})
+    model_name = get_config_value(model_config, "name", "unsloth/phi-4-unsloth-bnb-4bit")
+    use_fast_tokenizer = get_config_value(model_config, "use_fast_tokenizer", True)
+    trust_remote_code = get_config_value(model_config, "trust_remote_code", True)
+    model_revision = get_config_value(config, "model_revision", "main")
+    
+    # Unsloth configuration
+    unsloth_config = get_config_value(config, "unsloth", {})
+    unsloth_enabled = get_config_value(unsloth_config, "enabled", True)
+    
+    # Tokenizer configuration
+    tokenizer_config = get_config_value(config, "tokenizer", {})
+    max_seq_length = min(
+        get_config_value(tokenizer_config, "max_seq_length", 2048),
+        4096  # Maximum supported by most models
+    )
+    add_eos_token = get_config_value(tokenizer_config, "add_eos_token", True)
+    chat_template = get_config_value(tokenizer_config, "chat_template", None)
+    padding_side = get_config_value(tokenizer_config, "padding_side", "right")
+    
+    log_info(f"Loading model: {model_name} (revision: {model_revision})")
+    log_info(f"Max sequence length: {max_seq_length}")
+    
     try:
-        if not unsloth_available:
-            logger.error("Unsloth is required for training with pre-quantized model")
-            logger.error("Please ensure unsloth is in requirements.txt")
-            raise ImportError("Unsloth is required for this training setup")
-        
-        # Get model name correctly from config
-        model_name = config.get("model_name") or config.get("model", {}).get("name")
-        logger.info(f"Loading model: {model_name}")
-        
-        if not model_name:
-            raise ValueError("Model name not found in configuration. Please check your transformers_config.json file.")
+        if unsloth_enabled and unsloth_available:
+            log_info("Using Unsloth for 4-bit quantized model and LoRA")
+            # Load using Unsloth
+            from unsloth import FastLanguageModel
+            model, tokenizer = FastLanguageModel.from_pretrained(
+                model_name=model_name,
+                max_seq_length=max_seq_length,
+                dtype=get_config_value(config, "torch_dtype", "bfloat16"),
+                revision=model_revision,
+                trust_remote_code=trust_remote_code,
+                use_flash_attention_2=get_config_value(config, "use_flash_attention", True)
+            )
             
-        logger.info("Using Unsloth optimizations with pre-quantized model")
-        
-        # First detect if we have a GPU
-        if torch.cuda.is_available():
-            gpu_count = torch.cuda.device_count()
-            logger.info(f"Found {gpu_count} CUDA devices")
-        else:
-            logger.warning("No CUDA devices detected. Training will be slow on CPU!")
-            gpu_count = 0
-        
-        # Set default dtype for better numerics
-        if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
-            # Use bfloat16 for Ampere or newer
-            dtype = torch.bfloat16
-            logger.info("Using bfloat16 precision (Ampere+ GPU)")
-        elif torch.cuda.is_available():
-            # Use float16 for older GPUs
-            dtype = torch.float16
-            logger.info("Using float16 precision (pre-Ampere GPU)")
+            # Configure tokenizer settings
+            tokenizer.padding_side = padding_side
+            if add_eos_token and tokenizer.eos_token is None:
+                log_info("Setting EOS token")
+                tokenizer.add_special_tokens({"eos_token": "</s>"})
+            
+            # Set chat template if specified
+            if chat_template:
+                log_info(f"Setting chat template: {chat_template}")
+                if hasattr(tokenizer, "chat_template"):
+                    tokenizer.chat_template = chat_template
+                else:
+                    log_info("Tokenizer does not support chat templates, using default formatting")
+            
+            # Apply LoRA
+            lora_r = get_config_value(unsloth_config, "r", 16)
+            lora_alpha = get_config_value(unsloth_config, "alpha", 32)
+            lora_dropout = get_config_value(unsloth_config, "dropout", 0)
+            target_modules = get_config_value(unsloth_config, "target_modules", 
+                ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
+            
+            log_info(f"Applying LoRA with r={lora_r}, alpha={lora_alpha}, dropout={lora_dropout}")
+            model = FastLanguageModel.get_peft_model(
+                model,
+                r=lora_r,
+                target_modules=target_modules,
+                lora_alpha=lora_alpha,
+                lora_dropout=lora_dropout,
+                bias="none",
+                use_gradient_checkpointing=get_config_value(config, "training.gradient_checkpointing", True),
+                random_state=0,
+                max_seq_length=max_seq_length,
+                modules_to_save=None
+            )
         else:
-            # CPU, use default dtype
-            dtype = None
-            logger.info("Using default precision (CPU)")
-        
-        # Check for flash attention as the last dependency check
-        use_flash_attention = config.get("use_flash_attention", True)
-        if use_flash_attention and not find_spec("flash_attn"):
-            logger.warning("flash-attn not found. Will continue without flash attention.")
-            logger.warning("To use flash attention, install with: pip install flash-attn --no-build-isolation")
-            use_flash_attention = False
-        
-        # Set device map based on config or default to "auto"
-        device_map = config.get("hardware", {}).get("hardware_setup", {}).get("device_map", "auto")
-        
-        # Calculate max memory settings if multiple GPUs are available
-        max_memory = None
-        if gpu_count > 1:
-            memory_per_gpu = config.get("hardware", {}).get("specs", {}).get("vram_per_gpu", 24)
-            max_memory = {i: f"{int(memory_per_gpu * 0.85)}GiB" for i in range(gpu_count)}
-            max_memory["cpu"] = "64GiB"  # Allow CPU offloading if needed
-        
-        # Load model with proper error handling for out-of-memory
-        try:
-            # Improved memory settings for multi-GPU setup
-            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+            # Standard HuggingFace loading
+            log_info("Using standard HuggingFace model loading (Unsloth not available or disabled)")
+            from transformers import AutoModelForCausalLM, AutoTokenizer
             
-            model, tokenizer = FastLanguageModel.from_pretrained(
-                model_name=model_name,
-                max_seq_length=config.get("max_seq_length", 2048) or config.get("tokenizer", {}).get("max_seq_length", 2048),
-                dtype=dtype,
-                device_map=device_map,
-                max_memory=max_memory,
-                # Don't explicitly use flash attention config here, let Unsloth handle it
+            # Load tokenizer first
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                trust_remote_code=trust_remote_code,
+                use_fast=use_fast_tokenizer,
+                revision=model_revision,
+                padding_side=padding_side
             )
-        except RuntimeError as e:
-            if "CUDA out of memory" in str(e):
-                logger.error("Out of GPU memory. Consider using a smaller batch size or gradient accumulation steps.")
-                raise
-            else:
-                # Try again with CPU placement to see if it's a memory issue
-                logger.warning(f"Error loading model on default device: {str(e)}")
-                logger.warning("Attempting to load with device_map='cpu' and no specific dtype")
-                model, tokenizer = FastLanguageModel.from_pretrained(
-                    model_name=model_name,
-                    max_seq_length=config.get("max_seq_length", 2048) or config.get("tokenizer", {}).get("max_seq_length", 2048),
-                    dtype=None,
-                    device_map={"": "cpu"},
-                )
-                logger.warning("Model loaded on CPU. Training will be very slow.")
-        
-        # Ensure model and optimizer init is on the same device
-        logger.info(f"Model device map: {model.hf_device_map if hasattr(model, 'hf_device_map') else 'Not available'}")
             
-        # Apply Unsloth's training optimizations with config parameters
-        unsloth_config = config.get("unsloth", {})
-        
-        # Get dropout value; if not explicitly zero, warn about performance implications
-        lora_dropout = unsloth_config.get("dropout", 0.05)
-        if lora_dropout > 0:
-            logger.warning(f"Unsloth works best with dropout=0, but config has dropout={lora_dropout}")
-            logger.warning("This will impact performance but training will still work")
-            logger.warning("Consider setting dropout=0 in your config for better performance")
-        
-        # Apply optimizations
-        model = FastLanguageModel.get_peft_model(
-            model,
-            r=unsloth_config.get("r", 32),
-            target_modules=unsloth_config.get("target_modules", 
-                ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]),
-            lora_alpha=unsloth_config.get("alpha", 16),
-            lora_dropout=lora_dropout,  # Using the value from config or default
-            bias="none",
-            use_gradient_checkpointing=config.get("gradient_checkpointing", True) or config.get("training", {}).get("gradient_checkpointing", True),
-            random_state=config.get("seed", 42),
-        )
-        logger.info("Unsloth optimizations applied successfully")
-
-        # Set up tokenizer settings
-        chat_template = config.get("chat_template") or config.get("tokenizer", {}).get("chat_template")
-        if chat_template:
-            try:
-                # Get the correct chat template for phi models
-                template = get_chat_template("phi")
-                # Correctly apply the template to the tokenizer (it's a string)
-                if isinstance(template, str):
-                    tokenizer.chat_template = template
-                    logger.info("Set phi chat template (string)")
+            # Configure tokenizer settings
+            if add_eos_token and tokenizer.eos_token is None:
+                log_info("Setting EOS token")
+                tokenizer.add_special_tokens({"eos_token": "</s>"})
+            
+            # Set chat template if specified
+            if chat_template:
+                log_info(f"Setting chat template: {chat_template}")
+                if hasattr(tokenizer, "chat_template"):
+                    tokenizer.chat_template = chat_template
                 else:
-                    # If it's not a string, it's likely already a template object
-                    tokenizer.chat_template = template
-                    logger.info("Set phi chat template (object)")
-            except Exception as e:
-                logger.warning(f"Failed to set chat template: {str(e)}")
-                logger.warning("Chat formatting may not work correctly, but training can continue")
+                    log_info("Tokenizer does not support chat templates, using default formatting")
+            
+            # Now load model with updated tokenizer
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                trust_remote_code=trust_remote_code,
+                revision=model_revision,
+                torch_dtype=torch.bfloat16 if get_config_value(config, "torch_dtype", "bfloat16") == "bfloat16" else torch.float16,
+                device_map="auto" if CUDA_AVAILABLE else None
+            )
+            
+            # Apply PEFT/LoRA if enabled but using standard loading
+            if peft_available and get_config_value(unsloth_config, "enabled", True):
+                log_info("Applying standard PEFT/LoRA configuration")
+                from peft import LoraConfig, get_peft_model
+                
+                lora_r = get_config_value(unsloth_config, "r", 16)
+                lora_alpha = get_config_value(unsloth_config, "alpha", 32)
+                lora_dropout = get_config_value(unsloth_config, "dropout", 0)
+                target_modules = get_config_value(unsloth_config, "target_modules", 
+                    ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
+                
+                log_info(f"Applying LoRA with r={lora_r}, alpha={lora_alpha}, dropout={lora_dropout}")
+                lora_config = LoraConfig(
+                    r=lora_r,
+                    lora_alpha=lora_alpha,
+                    target_modules=target_modules,
+                    lora_dropout=lora_dropout,
+                    bias="none",
+                    task_type="CAUSAL_LM"
+                )
+                model = get_peft_model(model, lora_config)
         
-        # Ensure proper token settings
-        if tokenizer.pad_token_id is None:
-            tokenizer.pad_token_id = tokenizer.eos_token_id
-            logger.info(f"Set pad_token_id to eos_token_id: {tokenizer.pad_token_id}")
+        # Print model summary
+        log_info(f"Model loaded successfully: {model.__class__.__name__}")
+        if hasattr(model, "print_trainable_parameters"):
+            model.print_trainable_parameters()
+        else:
+            total_params = sum(p.numel() for p in model.parameters())
+            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+            log_info(f"Model has {total_params:,} parameters, {trainable_params:,} trainable ({trainable_params/total_params:.2%})")
         
         return model, tokenizer
-    
+        
     except Exception as e:
-        logger.error(f"Error in model/tokenizer loading: {str(e)}")
-        logger.error("If missing dependencies, check the requirements.txt file")
-        raise
+        log_info(f"Error loading model: {str(e)}")
+        traceback.print_exc()
+        return None, None
 
-def load_dataset_with_mapping(dataset_config):
-    """Load dataset and apply appropriate column mappings."""
+def load_dataset_with_mapping(config):
+    """
+    Load dataset from Hugging Face or local files and apply necessary transformations.
+    
+    Args:
+        config (dict): Dataset configuration dictionary
+    
+    Returns:
+        Dataset: The loaded and processed dataset
+    """
+    # Extract dataset configuration
+    dataset_info = get_config_value(config, "dataset", {})
+    dataset_name = get_config_value(dataset_info, "name", None)
+    dataset_split = get_config_value(dataset_info, "split", "train")
+    
+    # Data formatting configuration
+    formatting_config = get_config_value(config, "data_formatting", {})
+    
+    if not dataset_name:
+        raise ValueError("Dataset name not specified in config")
+    
+    log_info(f"Loading dataset: {dataset_name} (split: {dataset_split})")
+    
     try:
-        # Load dataset
-        dataset_name = dataset_config.get("dataset", {}).get("name", "")
-        dataset_split = dataset_config.get("dataset", {}).get("split", "train")
+        # Load dataset from Hugging Face or local path
+        from datasets import load_dataset
         
-        if not dataset_name:
-            raise ValueError("Dataset name not provided in configuration")
+        # Check if it's a local path or Hugging Face dataset
+        if os.path.exists(dataset_name) or os.path.exists(os.path.join(os.getcwd(), dataset_name)):
+            log_info(f"Loading dataset from local path: {dataset_name}")
+            # Local dataset - check if it's a directory or file
+            if os.path.isdir(dataset_name):
+                # Directory - look for data files
+                dataset = load_dataset(
+                    "json", 
+                    data_files={"train": os.path.join(dataset_name, "*.json")},
+                    split=dataset_split
+                )
+            else:
+                # Single file
+                dataset = load_dataset(
+                    "json",
+                    data_files={"train": dataset_name},
+                    split=dataset_split
+                )
+        else:
+            # Hugging Face dataset
+            log_info(f"Loading dataset from Hugging Face: {dataset_name}")
+            dataset = load_dataset(dataset_name, split=dataset_split)
         
-        logger.info(f"Loading pre-processed dataset {dataset_name}, split {dataset_split}")
+        log_info(f"Dataset loaded with {len(dataset)} examples")
         
-        try:
-            dataset = load_dataset(dataset_name, split=dataset_split)
+        # Check if dataset contains required fields
+        required_fields = ["conversations"]
+        missing_fields = [field for field in required_fields if field not in dataset.column_names]
+        
+        if missing_fields:
+            log_info(f"WARNING: Dataset missing required fields: {missing_fields}")
+            log_info("Attempting to map dataset structure to required format")
             
-            # Verify the dataset was actually loaded and is not None
-            if dataset is None:
-                raise ValueError(f"Dataset {dataset_name} (split {dataset_split}) loaded as None - check dataset exists and is accessible")
-                
-            # Check if the dataset is empty
-            if len(dataset) == 0:
-                raise ValueError(f"Dataset {dataset_name} (split {dataset_split}) is empty (contains 0 examples)")
-                
-            # Verify conversations field specifically
-            if "conversations" not in dataset.column_names:
-                raise ValueError(f"Dataset {dataset_name} missing required 'conversations' column")
-                
-            # Validate conversation structure
-            if len(dataset) > 0:
-                sample = dataset[0]
-                conversations = sample.get("conversations", [])
-                
-                if conversations:
-                    first_conv = conversations[0]
-                    if isinstance(first_conv, dict):
-                        # Check actual fields
-                        fields = list(first_conv.keys())
-                        logger.info(f"Conversation fields: {fields}")
-                        
-                        # Verify only 'content' field exists
-                        if fields == ["content"]:
-                            logger.info("Confirmed conversations have correct format with only 'content' field")
-                        else:
-                            logger.warning(f"Unexpected conversation fields: {fields}")
-                            logger.warning("Expected only 'content' field")
-                
-            # Check a sample of conversation entries to validate structure
-            logger.info("Validating conversation structure...")
-            for i in range(min(5, len(dataset))):
-                conv = dataset[i].get("conversations")
-                if conv is None:
-                    logger.warning(f"Example {i} has None as 'conversations' value")
-                elif not isinstance(conv, list):
-                    logger.warning(f"Example {i} has non-list 'conversations': {type(conv)}")
-                elif len(conv) == 0:
-                    logger.warning(f"Example {i} has empty conversations list")
+            # Implement conversion logic based on dataset structure
+            if "messages" in dataset.column_names:
+                log_info("Converting 'messages' field to 'conversations' format")
+                dataset = dataset.map(
+                    lambda x: {"conversations": x["messages"]},
+                    remove_columns=["messages"]
+                )
+            elif "text" in dataset.column_names:
+                log_info("Converting plain text to conversations format")
+                dataset = dataset.map(
+                    lambda x: {"conversations": [{"role": "user", "content": x["text"]}]},
+                    remove_columns=["text"]
+                )
+            else:
+                raise ValueError(f"Cannot convert dataset format - missing required fields and no conversion path available")
+        
+        # Log dataset info
+        log_info(f"Dataset has {len(dataset)} examples and columns: {dataset.column_names}")
+        
+        # Show a few examples for verification
+        for i in range(min(3, len(dataset))):
+            example = dataset[i]
+            log_info(f"Example {i}:")
+            for key, value in example.items():
+                if key == "conversations":
+                    log_info(f"  conversations: {len(value)} messages")
+                    # Show first message only to avoid cluttering logs
+                    if value and len(value) > 0:
+                        first_msg = value[0]
+                        if isinstance(first_msg, dict) and "content" in first_msg:
+                            content = first_msg["content"]
+                            log_info(f"  First message: {content[:50]}..." if len(content) > 50 else f"  First message: {content}")
                 else:
-                    # Look at the first conversation entry
-                    first_entry = conv[0]
-                    if isinstance(first_entry, dict) and "content" in first_entry:
-                        logger.info(f"Content field example: {str(first_entry['content'])[:50]}...")
-                    else:
-                        logger.warning(f"Example {i} missing 'content' key in conversation")
-                
-        except Exception as dataset_error:
-            logger.error(f"Failed to load dataset {dataset_name}: {str(dataset_error)}")
-            logger.error("Make sure the dataset exists and you have proper access permissions")
-            logger.error("This could be due to authentication issues with your HF_TOKEN")
-            raise
+                    log_info(f"  {key}: {value}")
         
         return dataset
         
     except Exception as e:
-        logger.error(f"Error loading dataset: {str(e)}")
-        return 1
+        log_info(f"Error loading dataset: {str(e)}")
+        traceback.print_exc()
+        return None
 
 def format_phi_chat(messages, dataset_config):
     """Format messages according to phi-4's chat template and dataset config.
@@ -528,110 +592,292 @@ class SimpleDataCollator:
             # Return empty batch if no valid examples
             return {k: [] for k in batch}
 
+def log_gpu_memory_usage(step=None, frequency=50, clear_cache_threshold=0.9, label=None):
+    """
+    Log GPU memory usage statistics with optional cache clearing
+    
+    Args:
+        step: Current training step (if None, logs regardless of frequency)
+        frequency: How often to log when step is provided
+        clear_cache_threshold: Fraction of memory used that triggers cache clearing (0-1)
+        label: Optional label for the log message (e.g., "Initial", "Error", "Step")
+    """
+    if not CUDA_AVAILABLE:
+        return
+    
+    # Only log every 'frequency' steps if step is provided
+    if step is not None and frequency > 0 and step % frequency != 0:
+        return
+        
+    # Get memory usage for each GPU
+    memory_info = []
+    for i in range(NUM_GPUS):
+        allocated = torch.cuda.memory_allocated(i) / (1024 ** 2)  # MB
+        reserved = torch.cuda.memory_reserved(i) / (1024 ** 2)    # MB
+        max_mem = torch.cuda.max_memory_allocated(i) / (1024 ** 2) # MB
+        
+        # Calculate percentage of reserved memory that's allocated
+        usage_percent = (allocated / reserved) * 100 if reserved > 0 else 0
+        memory_info.append(f"GPU {i}: {allocated:.1f}MB/{reserved:.1f}MB ({usage_percent:.1f}%, max: {max_mem:.1f}MB)")
+        
+        # Automatically clear cache if over threshold
+        if clear_cache_threshold > 0 and reserved > 0 and (allocated / reserved) > clear_cache_threshold:
+            log_info(f"Clearing CUDA cache for GPU {i} - high utilization ({allocated:.1f}/{reserved:.1f}MB)")
+            with torch.cuda.device(i):
+                torch.cuda.empty_cache()
+    
+    prefix = f"{label} " if label else ""
+    log_info(f"{prefix}GPU Memory: {', '.join(memory_info)}")
+
 class LoggingCallback(TrainerCallback):
     def __init__(self, model=None, dataset=None):
         super().__init__()
         self.training_started = time.time()
         self.last_log_time = time.time()
-        self.last_step = 0
+        self.last_step_time = None
+        self.step_durations = []
+        self.best_loss = float('inf')
         self.model = model
         self.dataset = dataset
         
     def on_train_begin(self, args, state, control, **kwargs):
-        log_info(f"=== Training started at {time.strftime('%Y-%m-%d %H:%M:%S')} ===")
-        
-        # Log model info if available
-        if self.model is not None:
-            log_info(f"Model parameters: {sum(p.numel() for p in self.model.parameters())/1e6:.2f}M")
+        """Called at the beginning of training"""
+        try:
+            log_info(f"=== Training started at {time.strftime('%Y-%m-%d %H:%M:%S')} ===")
+            
+            # Log model info if available
+            if self.model is not None:
+                total_params = sum(p.numel() for p in self.model.parameters())
+                trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+                log_info(f"Model parameters: {total_params/1e6:.2f}M total, {trainable_params/1e6:.2f}M trainable")
+                
+            # Log dataset info if available
+            if self.dataset is not None:
+                log_info(f"Dataset size: {len(self.dataset)} examples")
+            
+            # Log important training parameters for visibility
+            total_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps * NUM_GPUS
+            total_steps = int(len(self.dataset or []) / (args.per_device_train_batch_size * NUM_GPUS * args.gradient_accumulation_steps) * args.num_train_epochs)
+            log_info(f"Training plan: {len(self.dataset or [])} examples over {args.num_train_epochs} epochs ≈ {total_steps} steps")
+            log_info(f"Batch size: {args.per_device_train_batch_size} × {args.gradient_accumulation_steps} steps × {NUM_GPUS} GPUs = {total_batch_size} total")
+            
+            # Log initial GPU memory usage with label
+            log_gpu_memory_usage(label="Initial")
+        except Exception as e:
+            logger.warning(f"Error logging training begin statistics: {str(e)}")
+    
+    def on_step_end(self, args, state, control, **kwargs):
+        """Called at the end of each step"""
+        try:
+            if state.global_step == 1 or state.global_step % args.logging_steps == 0:
+                # Track step timing
+                current_time = time.time()
+                if self.last_step_time:
+                    step_duration = current_time - self.last_step_time
+                    self.step_durations.append(step_duration)
+                    # Keep only last 100 steps for averaging
+                    if len(self.step_durations) > 100:
+                        self.step_durations.pop(0)
+                    avg_step_time = sum(self.step_durations) / len(self.step_durations)
+                    log_info(f"Step {state.global_step}: {step_duration:.2f}s (avg: {avg_step_time:.2f}s)")
+                
+                self.last_step_time = current_time
+                
+                # Log GPU memory usage with step number
+                log_gpu_memory_usage(state.global_step, args.logging_steps)
+                
+                # Log loss
+                if state.log_history:
+                    latest_logs = state.log_history[-1] if state.log_history else {}
+                    if "loss" in latest_logs:
+                        loss = latest_logs["loss"]
+                        log_info(f"Step {state.global_step} loss: {loss:.4f}")
+                        
+                        # Track best loss
+                        if loss < self.best_loss:
+                            self.best_loss = loss
+                            log_info(f"New best loss: {loss:.4f}")
+        except Exception as e:
+            logger.warning(f"Error logging step end statistics: {str(e)}")
+    
+    def on_train_end(self, args, state, control, **kwargs):
+        """Called at the end of training"""
+        try:
+            # Calculate training duration
+            training_time = time.time() - self.training_started
+            hours, remainder = divmod(training_time, 3600)
+            minutes, seconds = divmod(remainder, 60)
+            
+            log_info(f"=== Training completed at {time.strftime('%Y-%m-%d %H:%M:%S')} ===")
+            log_info(f"Training duration: {int(hours)}h {int(minutes)}m {int(seconds)}s")
+            log_info(f"Final step: {state.global_step}")
+            log_info(f"Best loss: {self.best_loss:.4f}")
+            
+            # Log final GPU memory usage
+            log_gpu_memory_usage(label="Final")
+        except Exception as e:
+            logger.warning(f"Error logging training end statistics: {str(e)}")
             
-        # Log dataset info if available
-        if self.dataset is not None:
-            log_info(f"Dataset size: {len(self.dataset)} examples")
+    # Other callback methods with proper error handling
+    def on_save(self, args, state, control, **kwargs):
+        """Called when a checkpoint is saved"""
+        try:
+            log_info(f"Saving checkpoint at step {state.global_step}")
+        except Exception as e:
+            logger.warning(f"Error in on_save: {str(e)}")
+    
+    def on_log(self, args, state, control, **kwargs):
+        """Called when a log is created"""
+        pass
         
-        # Log important training parameters for visibility
-        total_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps * NUM_GPUS
-        total_steps = int(len(self.dataset or []) / (args.per_device_train_batch_size * NUM_GPUS * args.gradient_accumulation_steps) * args.num_train_epochs)
-        log_info(f"Training plan: {len(self.dataset or [])} examples over {args.num_train_epochs} epochs ≈ {total_steps} steps")
-        log_info(f"Batch size: {args.per_device_train_batch_size} × {args.gradient_accumulation_steps} steps × {NUM_GPUS} GPUs = {total_batch_size} total")
+    def on_evaluate(self, args, state, control, **kwargs):
+        """Called when evaluation is performed"""
+        pass
         
-        # Log memory information in compact format
-        if CUDA_AVAILABLE:
-            memory_info = []
-            for i in range(NUM_GPUS):
-                allocated = torch.cuda.memory_allocated(i) / 1024**2
-                max_mem = torch.cuda.max_memory_allocated(i) / 1024**2
-                memory_info.append(f"GPU {i}: {allocated:.1f}MB (max: {max_mem:.1f}MB)")
-            log_info(f"Initial memory usage - {', '.join(memory_info)}")
+    # Only implement the methods we actually need, remove the others
+    def on_prediction_step(self, args, state, control, **kwargs):
+        """Called when prediction is performed"""
+        pass
+    
+    def on_save_model(self, args, state, control, **kwargs):
+        """Called when model is saved"""
+        try:
+            # Log memory usage after saving
+            log_gpu_memory_usage(label=f"Save at step {state.global_step}")
+        except Exception as e:
+            logger.warning(f"Error in on_save_model: {str(e)}")
+    
+    def on_epoch_end(self, args, state, control, **kwargs):
+        """Called at the end of an epoch"""
+        try:
+            epoch = state.epoch
+            log_info(f"Completed epoch {epoch:.2f}")
+            log_gpu_memory_usage(label=f"Epoch {epoch:.2f}")
+        except Exception as e:
+            logger.warning(f"Error in on_epoch_end: {str(e)}")
+    
+    def on_step_begin(self, args, state, control, **kwargs):
+        """Called at the beginning of a step"""
+        pass
 
 def check_dependencies():
-    """Check if all required dependencies are installed and in the correct order."""
-    missing_packages = []
-    order_issues = []
-    
-    # Define required packages with versions
+    """
+    Check for required and optional dependencies, ensuring proper versions and import order.
+    Returns True if all required dependencies are present, False otherwise.
+    """
+    # Define required packages with versions and descriptions
     required_packages = {
-        "unsloth": ">=2024.3",
-        "transformers": ">=4.38.0",
-        "peft": ">=0.9.0",
-        "accelerate": ">=0.27.0"
+        "unsloth": {"version": ">=2024.3", "feature": "fast 4-bit quantization and LoRA"},
+        "transformers": {"version": ">=4.38.0", "feature": "core model functionality"},
+        "peft": {"version": ">=0.9.0", "feature": "parameter-efficient fine-tuning"},
+        "accelerate": {"version": ">=0.27.0", "feature": "multi-GPU training"}
     }
     
-    # Check for required packages
-    for package, version in required_packages.items():
+    # Optional packages that enhance functionality
+    optional_packages = {
+        "flash_attn": {"feature": "faster attention computation"},
+        "bitsandbytes": {"feature": "quantization support"},
+        "optimum": {"feature": "model optimization"},
+        "wandb": {"feature": "experiment tracking"}
+    }
+    
+    # Store results
+    missing_packages = []
+    package_versions = {}
+    order_issues = []
+    
+    # Check required packages
+    log_info("Checking required dependencies...")
+    for package, info in required_packages.items():
+        version_req = info["version"]
+        feature = info["feature"]
+        
         try:
+            # Special handling for packages we've already checked
             if package == "unsloth" and not unsloth_available:
-                missing_packages.append(f"{package}{version}")
+                missing_packages.append(f"{package}{version_req}")
+                log_info(f"❌ {package} - {feature} MISSING")
+                continue
             elif package == "peft" and not peft_available:
-                missing_packages.append(f"{package}{version}")
-            else:
-                module = __import__(package)
-                logger.info(f"Using {package} version {getattr(module, '__version__', 'unknown')}")
+                missing_packages.append(f"{package}{version_req}")
+                log_info(f"❌ {package} - {feature} MISSING")
+                continue
+                
+            # Try to import and get version
+            module = __import__(package)
+            version = getattr(module, "__version__", "unknown")
+            package_versions[package] = version
+            log_info(f"✅ {package} v{version} - {feature}")
+            
         except ImportError:
-            missing_packages.append(f"{package}{version}")
+            missing_packages.append(f"{package}{version_req}")
+            log_info(f"❌ {package} - {feature} MISSING")
     
-    # Check import order
-    try:
-        import sys
-        modules = list(sys.modules.keys())
-        
-        if 'transformers' in modules and 'unsloth' in modules:
-            try:
-                transformers_idx = modules.index('transformers')
-                unsloth_idx = modules.index('unsloth')
-                if transformers_idx < unsloth_idx:
-                    order_issues.append("For optimal performance, unsloth should be imported before transformers")
-            except ValueError:
-                pass
-    except Exception as e:
-        logger.warning(f"Could not check module import order: {str(e)}")
-    
-    # Check optional dependencies
-    optional_packages = {
-        "flash_attn": "Flash attention support",
-        "bitsandbytes": "4-bit quantization support"
-    }
+    # Check optional packages
+    log_info("\nChecking optional dependencies...")
+    for package, info in optional_packages.items():
+        feature = info["feature"]
+        try:
+            __import__(package)
+            log_info(f"✅ {package} - {feature} available")
+        except ImportError:
+            log_info(f"⚠️ {package} - {feature} not available")
     
-    for package, feature in optional_packages.items():
-        if find_spec(package):
-            logger.info(f"Found {package} - {feature} enabled")
-        else:
-            logger.warning(f"{package} not found - {feature} will not be available")
+    # Check import order for optimal performance
+    if "transformers" in package_versions and "unsloth" in package_versions:
+        try:
+            import sys
+            modules = list(sys.modules.keys())
+            transformers_idx = modules.index("transformers")
+            unsloth_idx = modules.index("unsloth")
+            
+            if transformers_idx < unsloth_idx:
+                order_issue = "⚠️ For optimal performance, import unsloth before transformers"
+                order_issues.append(order_issue)
+                log_info(order_issue)
+            else:
+                log_info("✅ Import order: unsloth before transformers (optimal)")
+        except (ValueError, IndexError) as e:
+            log_info(f"⚠️ Could not verify import order: {str(e)}")
     
     # Report missing required packages
     if missing_packages:
-        logger.error("Critical dependencies missing:")
+        log_info("\n❌ Critical dependencies missing:")
         for pkg in missing_packages:
-            logger.error(f"  - {pkg}")
-        logger.error("Please install the missing dependencies with:")
-        logger.error(f"  pip install {' '.join(missing_packages)}")
+            log_info(f"  - {pkg}")
+        log_info("Please install missing dependencies with:")
+        log_info(f"  pip install {' '.join(missing_packages)}")
         return False
     
-    # Report order issues as warnings
-    for issue in order_issues:
-        logger.warning(issue)
-    
+    log_info("\n✅ All required dependencies satisfied!")
     return True
 
+def get_config_value(config, path, default=None):
+    """
+    Safely get a nested value from a config dictionary using a dot-separated path.
+    
+    Args:
+        config: The configuration dictionary
+        path: Dot-separated path to the value (e.g., "training.optimizer.lr")
+        default: Default value to return if path doesn't exist
+        
+    Returns:
+        The value at the specified path or the default value
+    """
+    if not config:
+        return default
+        
+    parts = path.split('.')
+    current = config
+    
+    for part in parts:
+        if isinstance(current, dict) and part in current:
+            current = current[part]
+        else:
+            return default
+            
+    return current
+
 def update_huggingface_space():
     """Update the Hugging Face Space with the current code."""
     log_info("Updating Hugging Face Space...")
@@ -709,381 +955,338 @@ def validate_huggingface_credentials():
         logger.warning(f"Error validating Hugging Face credentials: {str(e)}")
         return False
 
+def setup_environment(args):
+    """
+    Set up the training environment including logging, seed, and configurations.
+    
+    Args:
+        args: Command line arguments
+    
+    Returns:
+        tuple: (transformers_config, seed) - The loaded configuration and random seed
+    """
+    # Load environment variables first
+    load_env_variables()
+    
+    # Set random seed for reproducibility
+    seed = args.seed if args.seed is not None else int(time.time()) % 10000
+    set_seed(seed)
+    log_info(f"Using random seed: {seed}")
+    
+    # Load configuration
+    base_path = os.path.dirname(os.path.abspath(__file__))
+    config_file = args.config_file or os.path.join(base_path, "transformers_config.json")
+    
+    if not os.path.exists(config_file):
+        raise FileNotFoundError(f"Config file not found: {config_file}")
+        
+    log_info(f"Loading configuration from {config_file}")
+    transformers_config = load_configs(config_file)
+    
+    # Set up hardware environment variables if CUDA is available
+    if CUDA_AVAILABLE:
+        memory_fraction = get_config_value(transformers_config, "hardware.system_settings.cuda_memory_fraction", 0.75)
+        if memory_fraction < 1.0:
+            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = f"max_split_size_mb:128,expandable_segments:True"
+            log_info(f"Set CUDA memory allocation limit to expandable with max_split_size_mb:128")
+    
+    # Check dependencies before proceeding
+    if not check_dependencies():
+        raise RuntimeError("Critical dependencies missing")
+    
+    return transformers_config, seed
+
+def setup_model_and_tokenizer(config):
+    """
+    Load and configure the model and tokenizer.
+    
+    Args:
+        config: Complete configuration dictionary
+    
+    Returns:
+        tuple: (model, tokenizer) - The loaded model and tokenizer
+    """
+    log_info("Loading model and tokenizer...")
+    model, tokenizer = load_model_and_tokenizer(config)
+    
+    if model is None or tokenizer is None:
+        raise ValueError("Failed to load model or tokenizer")
+        
+    log_info(f"Model loaded successfully: {model.__class__.__name__}")
+    log_info(f"Tokenizer loaded: {tokenizer.__class__.__name__} (vocab size: {tokenizer.vocab_size})")
+    
+    return model, tokenizer
+
+def setup_dataset_and_collator(config, tokenizer):
+    """
+    Load and configure the dataset and data collator.
+    
+    Args:
+        config: Complete configuration dictionary
+        tokenizer: The tokenizer for the data collator
+    
+    Returns:
+        tuple: (dataset, data_collator) - The loaded dataset and configured data collator
+    """
+    dataset_config = get_config_value(config, "dataset", {})
+    
+    log_info("Loading dataset...")
+    dataset = load_dataset_with_mapping(dataset_config)
+    
+    # Validate dataset
+    if dataset is None:
+        raise ValueError("Dataset is None! Cannot proceed with training.")
+        
+    if not hasattr(dataset, '__len__') or len(dataset) == 0:
+        raise ValueError("Dataset is empty! Cannot proceed with training.")
+        
+    log_info(f"Dataset loaded with {len(dataset)} examples")
+    
+    # Create data collator
+    data_collator = SimpleDataCollator(tokenizer, dataset_config)
+    
+    return dataset, data_collator
+
+def create_training_arguments(config, dataset):
+    """
+    Create and configure training arguments for the Trainer.
+    
+    Args:
+        config: Complete configuration dictionary
+        dataset: The dataset to determine total steps
+    
+    Returns:
+        TrainingArguments: Configured training arguments
+    """
+    # Extract configuration sections
+    training_config = get_config_value(config, "training", {})
+    hardware_config = get_config_value(config, "hardware", {})
+    huggingface_config = get_config_value(config, "huggingface_hub", {})
+    distributed_config = get_config_value(config, "distributed_training", {})
+    
+    # Extract key training parameters
+    per_device_batch_size = get_config_value(training_config, "per_device_train_batch_size", 4)
+    gradient_accumulation_steps = get_config_value(training_config, "gradient_accumulation_steps", 8)
+    learning_rate = get_config_value(training_config, "learning_rate", 2e-5)
+    num_train_epochs = get_config_value(training_config, "num_train_epochs", 3)
+    
+    # Extract hardware settings
+    dataloader_workers = get_config_value(hardware_config, "system_settings.dataloader_num_workers", 
+                                     get_config_value(distributed_config, "dataloader_num_workers", 2))
+    pin_memory = get_config_value(hardware_config, "system_settings.dataloader_pin_memory", True)
+    
+    # BF16/FP16 settings - ensure only one is enabled
+    use_bf16 = get_config_value(training_config, "bf16", False)
+    use_fp16 = get_config_value(training_config, "fp16", False) if not use_bf16 else False
+    
+    # Configure distributed training
+    fsdp_config = get_config_value(distributed_config, "fsdp_config", {})
+    fsdp_enabled = get_config_value(fsdp_config, "enabled", False)
+    
+    ddp_config = get_config_value(distributed_config, "ddp_config", {})
+    ddp_find_unused_parameters = get_config_value(ddp_config, "find_unused_parameters", False)
+    
+    # Set up FSDP args if enabled
+    fsdp_args = None
+    if fsdp_enabled and NUM_GPUS > 1:
+        from accelerate import FullyShardedDataParallelPlugin
+        from torch.distributed.fsdp.fully_sharded_data_parallel import (
+            FullOptimStateDictConfig, FullStateDictConfig
+        )
+        
+        fsdp_plugin = FullyShardedDataParallelPlugin(
+            sharding_strategy=get_config_value(fsdp_config, "sharding_strategy", "FULL_SHARD"),
+            mixed_precision_policy=get_config_value(fsdp_config, "mixed_precision", "BF16"),
+            state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+            optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
+        )
+        
+        fsdp_args = {
+            "fsdp": fsdp_plugin,
+            "fsdp_transformer_layer_cls_to_wrap": ["LlamaDecoderLayer", "PhiDecoderLayer"]
+        }
+    
+    # Create and return training arguments
+    training_args = TrainingArguments(
+        output_dir=get_config_value(config, "checkpointing.output_dir", "./results"),
+        overwrite_output_dir=True,
+        num_train_epochs=num_train_epochs,
+        per_device_train_batch_size=per_device_batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        learning_rate=learning_rate,
+        weight_decay=get_config_value(training_config, "weight_decay", 0.01),
+        max_grad_norm=get_config_value(training_config, "max_grad_norm", 1.0),
+        warmup_ratio=get_config_value(training_config, "warmup_ratio", 0.03),
+        lr_scheduler_type=get_config_value(training_config, "lr_scheduler_type", "cosine"),
+        logging_steps=get_config_value(training_config, "logging_steps", 10),
+        save_strategy=get_config_value(config, "checkpointing.save_strategy", "steps"),
+        save_steps=get_config_value(config, "checkpointing.save_steps", 500),
+        save_total_limit=get_config_value(config, "checkpointing.save_total_limit", 3),
+        bf16=use_bf16,
+        fp16=use_fp16,
+        push_to_hub=get_config_value(huggingface_config, "push_to_hub", False),
+        hub_model_id=get_config_value(huggingface_config, "hub_model_id", None),
+        hub_strategy=get_config_value(huggingface_config, "hub_strategy", "every_save"),
+        hub_private_repo=get_config_value(huggingface_config, "hub_private_repo", True),
+        gradient_checkpointing=get_config_value(training_config, "gradient_checkpointing", True),
+        dataloader_pin_memory=pin_memory,
+        optim=get_config_value(training_config, "optim", "adamw_torch"),
+        ddp_find_unused_parameters=ddp_find_unused_parameters,
+        dataloader_drop_last=False,
+        dataloader_num_workers=dataloader_workers,
+        no_cuda=False if CUDA_AVAILABLE else True,
+        **({} if fsdp_args is None else fsdp_args)
+    )
+
+    log_info("Training arguments created successfully")
+    return training_args
+
+def configure_custom_dataloader(trainer, dataset, config, training_args):
+    """
+    Configure a custom dataloader for the trainer if needed.
+    
+    Args:
+        trainer: The Trainer instance to configure
+        dataset: The dataset to use
+        config: Complete configuration dictionary
+        training_args: The training arguments
+        
+    Returns:
+        None (modifies trainer in-place)
+    """
+    dataset_config = get_config_value(config, "dataset", {})
+    
+    # Check if we need a custom dataloader
+    if get_config_value(dataset_config, "data_loading.sequential_processing", True):
+        log_info("Using custom sequential dataloader")
+        
+        # Create sequential sampler to maintain dataset order
+        sequential_sampler = torch.utils.data.SequentialSampler(dataset)
+        log_info("Sequential sampler created")
+        
+        # Define custom dataloader getter
+        def custom_get_train_dataloader():
+            """Create a custom dataloader that maintains dataset order"""
+            # Get configuration values
+            batch_size = training_args.per_device_train_batch_size
+            drop_last = get_config_value(dataset_config, "data_loading.drop_last", False)
+            num_workers = training_args.dataloader_num_workers
+            pin_memory = training_args.dataloader_pin_memory
+            prefetch_factor = get_config_value(dataset_config, "data_loading.prefetch_factor", 2)
+            persistent_workers = get_config_value(dataset_config, "data_loading.persistent_workers", False)
+            
+            # Create DataLoader with sequential sampler
+            return DataLoader(
+                dataset,
+                batch_size=batch_size,
+                sampler=sequential_sampler,
+                collate_fn=trainer.data_collator,
+                drop_last=drop_last,
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                prefetch_factor=prefetch_factor if num_workers > 0 else None,
+                persistent_workers=persistent_workers if num_workers > 0 else False,
+            )
+        
+        # Override the default dataloader
+        trainer.get_train_dataloader = custom_get_train_dataloader
+
+def run_training(trainer, tokenizer, training_args):
+    """
+    Run the training process and handle model saving.
+    
+    Args:
+        trainer: Configured Trainer instance
+        tokenizer: The tokenizer to save with the model
+        training_args: Training arguments
+        
+    Returns:
+        int: 0 for success, 1 for failure
+    """
+    log_info("Starting training...")
+    trainer.train()
+    
+    log_info("Training complete! Saving final model...")
+    trainer.save_model()
+    tokenizer.save_pretrained(training_args.output_dir)
+    
+    # Push to Hub if configured
+    if training_args.push_to_hub:
+        log_info(f"Pushing model to Hugging Face Hub: {training_args.hub_model_id}")
+        trainer.push_to_hub()
+        
+    log_info("Training completed successfully!")
+    return 0
+
 def main():
+    """
+    Main entry point for the training script.
+    
+    Returns:
+        int: 0 for success, non-zero for failure
+    """
     # Set up logging
     logger.info("Starting training process")
     
     try:
-        # Check dependencies first, before any other operations
-        if not check_dependencies():
-            logger.error("Aborting due to missing critical dependencies")
-            return 1
-        
-        # Parse arguments
+        # Parse command line arguments
         args = parse_args()
         
-        # Load environment variables
-        load_env_variables()
+        # Set up environment and load configuration
+        transformers_config, seed = setup_environment(args)
         
-        # Validate Hugging Face credentials if we're going to use them
-        validate_huggingface_credentials()
-        
-        # Load configuration
+        # Load model and tokenizer
         try:
-            transformers_config = load_configs(args.config)
-            hardware_config = transformers_config.get("hardware", {})
-            dataset_config = transformers_config.get("dataset", {})
-            logger.info("Configuration loaded successfully")
+            model, tokenizer = setup_model_and_tokenizer(transformers_config)
         except Exception as e:
-            logger.error(f"Error loading configuration: {e}")
+            logger.error(f"Error setting up model: {str(e)}")
             return 1
         
-        # Check if we're in distributed mode
-        is_distributed = "WORLD_SIZE" in os.environ and int(os.environ.get("WORLD_SIZE", "1")) > 1
-        if is_distributed:
-            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
-            log_info(f"Running in distributed mode with {os.environ.get('WORLD_SIZE')} processes, local_rank: {local_rank}")
-        else:
-            log_info("Running in non-distributed mode (single process)")
-        
-        # Set random seed for reproducibility
-        seed = transformers_config.get("seed", 42)
-        set_seed(seed)
-        logger.info(f"Set random seed to {seed}")
-        
-        # Load model and tokenizer using the consolidated config
-        model, tokenizer = load_model_and_tokenizer(transformers_config)
-        
-        # Empty CUDA cache to ensure clean state
-        if CUDA_AVAILABLE:
-            torch.cuda.empty_cache()
-            log_info("Cleared CUDA cache")
+        # Load dataset and create data collator
+        try:
+            dataset, data_collator = setup_dataset_and_collator(transformers_config, tokenizer)
+        except Exception as e:
+            logger.error(f"Error setting up dataset: {str(e)}")
+            return 1
         
-        # Setup environment variable for CUDA memory allocation
-        if CUDA_AVAILABLE:
-            system_settings = hardware_config.get("system_settings", {})
-            cuda_memory_fraction = system_settings.get("cuda_memory_fraction", 0.85)
+        # Configure training arguments
+        try:
+            training_args = create_training_arguments(transformers_config, dataset)
+        except Exception as e:
+            logger.error(f"Error configuring training arguments: {str(e)}")
+            return 1
             
-            if cuda_memory_fraction < 1.0:
-                os.environ["PYTORCH_CUDA_ALLOC_CONF"] = f"max_split_size_mb:128,expandable_segments:True"
-                log_info(f"Set CUDA memory allocation limit to expandable with max_split_size_mb:128")
+        # Initialize trainer with callbacks
+        log_info("Initializing Trainer")
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=dataset,
+            data_collator=data_collator,
+            callbacks=[LoggingCallback(model=model, dataset=dataset)],
+        )
         
+        # Configure custom dataloader if needed
         try:
-            log_info("Loading dataset...")
-            dataset = load_dataset_with_mapping(dataset_config)
-            
-            # Extra validation to catch None/empty dataset issues
-            if dataset is None:
-                logger.error("Dataset is None! Cannot proceed with training.")
-                return 1
-                
-            if not hasattr(dataset, '__len__') or len(dataset) == 0:
-                logger.error("Dataset is empty! Cannot proceed with training.")
-                return 1
-                
-            log_info(f"Dataset loaded with {len(dataset)} examples")
-            
-            # Minimal validation before proceeding
-            if dataset is None or len(dataset) == 0:
-                logger.error("Dataset is empty or None! Cannot proceed with training.")
-                return 1
-            
-            # Create data collator
-            data_collator = SimpleDataCollator(tokenizer, dataset_config)
-            
-            # Verify precision settings - ensure only one of bf16/fp16 is set, with bf16 taking precedence
-            # First check hardware config, then transformers config
-            use_bf16 = False
-            use_fp16 = False
-            
-            # Check hardware config first
-            hardware_precision = hardware_config.get("training_optimizations", {}).get("mixed_precision", "")
-            if hardware_precision.lower() == "bf16":
-                use_bf16 = True
-                log_info("Using BF16 precision from hardware config")
-            elif hardware_precision.lower() == "fp16":
-                use_fp16 = True
-                log_info("Using FP16 precision from hardware config")
-            else:
-                # Fall back to transformers config
-                use_bf16 = transformers_config.get("bf16", False) or transformers_config.get("torch_dtype", "") == "bfloat16"
-                use_fp16 = transformers_config.get("fp16", False) and not use_bf16  # Only use fp16 if bf16 is not set
-                log_info(f"Using precision: {'bf16' if use_bf16 else 'fp16' if use_fp16 else 'full precision'}")
-            
-            # Get per device batch size - from transformers config, but possibly overridden by hardware config
-            per_device_batch_size = transformers_config.get("training", {}).get("per_device_train_batch_size", 16)
-            gradient_accumulation_steps = transformers_config.get("training", {}).get("gradient_accumulation_steps", 3)
-            
-            # Get multi-GPU strategy from hardware config (default to data_parallel)
-            multi_gpu_strategy = hardware_config.get("training_optimizations", {}).get("multi_gpu_strategy", "data_parallel")
-            logger.info(f"Multi-GPU strategy: {multi_gpu_strategy}")
-            
-            # For multi-GPU setup, adjust for better balance
-            if CUDA_AVAILABLE and NUM_GPUS > 1:
-                log_info(f"Multi-GPU setup: Adjusting for {NUM_GPUS} GPUs")
-            
-            # Set up FSDP for multi-GPU training if specified and in distributed mode
-            fsdp_config = None
-            if multi_gpu_strategy == "fsdp" and is_distributed and NUM_GPUS > 1:
-                try:
-                    from torch.distributed.fsdp import (
-                        FullyShardedDataParallel as FSDP,
-                        MixedPrecision,
-                        BackwardPrefetch,
-                        ShardingStrategy,
-                        CPUOffload,
-                    )
-                    from torch.distributed.fsdp.wrap import (
-                        transformer_auto_wrap_policy,
-                        enable_wrap,
-                        wrap,
-                    )
-                    
-                    log_info("Using FSDP for distributed training")
-                    
-                    # Configure FSDP
-                    fsdp_config = {
-                        "fsdp_transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
-                        "fsdp_offload_params": False,
-                        "fsdp_backward_prefetch": "BACKWARD_PRE",
-                        "fsdp_min_num_params": 1e6,
-                        "fsdp_sharding_strategy": 1,  # FULL_SHARD
-                    }
-                    
-                    if use_bf16 or use_fp16:
-                        precision_type = "bf16" if use_bf16 else "fp16"
-                        fsdp_config["fsdp_state_dict_type"] = "FULL_STATE_DICT"
-                        log_info(f"FSDP using mixed precision: {precision_type}")
-                except ImportError:
-                    log_info("FSDP imports failed, falling back to standard DDP")
-                    fsdp_config = None
-            elif multi_gpu_strategy == "fsdp" and not is_distributed:
-                log_info("FSDP disabled: requires distributed environment (use torchrun or accelerate)")
-                log_info("Using DataParallel for multi-GPU training instead")
-            else:
-                log_info(f"Using {multi_gpu_strategy} for multi-GPU training")
-            
-            # Get system settings from hardware config
-            dataloader_workers = hardware_config.get("system_settings", {}).get("dataloader_num_workers", 2)
-            pin_memory = hardware_config.get("system_settings", {}).get("dataloader_pin_memory", True)
-            
-            # Set up training arguments
-            log_info("Setting up training arguments")
-            
-            # Handle FSDP configuration
-            fsdp_config = transformers_config.get("distributed_training", {}).get("fsdp_config", {})
-            fsdp_enabled = fsdp_config.get("enabled", False)
-            
-            # Only set FSDP args if explicitly enabled
-            fsdp_args = None
-            if fsdp_enabled and is_distributed and NUM_GPUS > 1:
-                fsdp_args = {
-                    "fsdp": ["full_shard", "auto_wrap"],
-                    "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
-                    "fsdp_offload_params": fsdp_config.get("offload_params", False),
-                    "fsdp_state_dict_type": "FULL_STATE_DICT",
-                    "fsdp_sharding_strategy": 1,  # FULL_SHARD
-                }
-                log_info("FSDP configuration enabled")
-            else:
-                log_info("FSDP disabled, using standard data parallel")
-
-            # Check if we're running in a Space
-            is_space = bool(os.environ.get("SPACE_ID"))
-            
-            # Create training arguments
-            training_args = TrainingArguments(
-                output_dir=transformers_config.get("output_dir", "./results") or transformers_config.get("checkpointing", {}).get("output_dir", "./results"),
-                num_train_epochs=transformers_config.get("training", {}).get("num_train_epochs", 3),
-                per_device_train_batch_size=per_device_batch_size,
-                gradient_accumulation_steps=gradient_accumulation_steps,
-                learning_rate=transformers_config.get("training", {}).get("learning_rate", 2e-5),
-                weight_decay=transformers_config.get("training", {}).get("weight_decay", 0.01),
-                warmup_ratio=transformers_config.get("training", {}).get("warmup_ratio", 0.05),
-                lr_scheduler_type=transformers_config.get("training", {}).get("lr_scheduler_type", "cosine"),
-                logging_steps=transformers_config.get("training", {}).get("logging_steps", 10),
-                save_strategy=transformers_config.get("checkpointing", {}).get("save_strategy", "steps"),
-                save_steps=transformers_config.get("checkpointing", {}).get("save_steps", 100),
-                save_total_limit=transformers_config.get("checkpointing", {}).get("save_total_limit", 3),
-                fp16=use_fp16,
-                bf16=use_bf16,
-                max_grad_norm=transformers_config.get("training", {}).get("max_grad_norm", 1.0),
-                push_to_hub=transformers_config.get("huggingface_hub", {}).get("push_to_hub", False),
-                hub_model_id=transformers_config.get("huggingface_hub", {}).get("hub_model_id", None),
-                hub_token=None if is_space else os.environ.get("HF_TOKEN", None),
-                report_to="tensorboard",
-                remove_unused_columns=False,  # Keep all columns
-                gradient_checkpointing=transformers_config.get("training", {}).get("gradient_checkpointing", True),
-                dataloader_pin_memory=pin_memory,
-                optim=transformers_config.get("training", {}).get("optim", "adamw_torch"),
-                ddp_find_unused_parameters=False,  # Improve distributed training efficiency
-                dataloader_drop_last=False,  # Process all examples
-                dataloader_num_workers=dataloader_workers,
-                no_cuda=False if CUDA_AVAILABLE else True,  # Use CUDA if available
-                **({} if fsdp_args is None else fsdp_args)  # Only include FSDP args if configured
-            )
-
-            log_info("Training arguments created successfully")
-            
-            # Validate dataset before creating sampler
-            if dataset is None:
-                raise ValueError("Dataset is None - cannot create sampler")
-            
-            # Create sequential sampler to maintain original dataset order
-            sequential_sampler = torch.utils.data.SequentialSampler(dataset)
-            log_info("Sequential sampler created")
-            
-            # Initialize trainer first
-            log_info("Initializing Trainer")
-            trainer = Trainer(
-                model=model,
-                args=training_args,
-                train_dataset=dataset,
-                data_collator=data_collator,
-                callbacks=[LoggingCallback(model=model, dataset=dataset)],
-            )
-            
-            # Then override the get_train_dataloader method
-            def custom_get_train_dataloader():
-                """Custom dataloader that preserves original dataset order"""
-                log_info("Creating sequential dataloader to maintain original dataset order")
-                
-                # Safety check - make sure dataset exists and is not None
-                if dataset is None:
-                    raise ValueError("Dataset is None - cannot create dataloader")
-                
-                # Make sure dataset is not empty
-                if len(dataset) == 0:
-                    raise ValueError("Dataset is empty - cannot create dataloader")
-                
-                # Create a simple sequential sampler
-                sequential_sampler = torch.utils.data.SequentialSampler(dataset)
-                
-                # Verification of sequence preservation flags - simplified
-                data_loading_config = dataset_config.get("data_loading", {})
-                shuffle_enabled = data_loading_config.get("shuffle", False)
-                
-                if shuffle_enabled:
-                    log_info("WARNING: Shuffle is enabled in configuration! This will be overridden to preserve order.")
-                    # We enforce sequential processing regardless of config
-                
-                # Log our approach clearly
-                log_info("Using SequentialSampler to guarantee dataset order is preserved based on prompt_number")
-                
-                # Verify column order and check for 'conversations' field
-                expected_order = ["prompt_number", "article_id", "conversations"]
-                if hasattr(dataset, 'column_names'):
-                    actual_order = dataset.column_names
-                    
-                    # Verify all required fields exist
-                    missing_fields = [field for field in ["conversations"] if field not in actual_order]
-                    if missing_fields:
-                        raise ValueError(f"Dataset missing critical fields: {missing_fields}")
-                    
-                    if actual_order == expected_order:
-                        log_info(f"Confirmed dataset columns are in expected order: {', '.join(expected_order)}")
-                    else:
-                        log_info(f"Note: Dataset columns ({', '.join(actual_order)}) are not in expected order ({', '.join(expected_order)})")
-                        log_info("This is handled correctly by field-based access, but noting for clarity")
-                
-                log_info("Dataset is pre-processed with prompt_number field indicating the correct sequence")
-                
-                # Validate a few samples before proceeding
-                for i in range(min(3, len(dataset))):
-                    sample = dataset[i]
-                    if "conversations" not in sample:
-                        log_info(f"WARNING: Sample {i} missing 'conversations' field")
-                    elif sample["conversations"] is None:
-                        log_info(f"WARNING: Sample {i} has None 'conversations' field")
-                    elif not isinstance(sample["conversations"], list):
-                        log_info(f"WARNING: Sample {i} has non-list 'conversations' field: {type(sample['conversations'])}")
-                
-                # Calculate batch size based on device availability
-                if getattr(training_args, "no_cuda", False):
-                    batch_size = training_args.per_device_train_batch_size
-                else:
-                    batch_size = max(training_args.per_device_train_batch_size * max(1, NUM_GPUS), 1)
-                
-                log_info(f"Using sequential sampler with batch size {batch_size}")
-                
-                # Return DataLoader with sequential sampler and extra error handling
-                try:
-                    return torch.utils.data.DataLoader(
-                        dataset,
-                        batch_size=batch_size,
-                        sampler=sequential_sampler,  # Always use sequential sampler
-                        collate_fn=data_collator,
-                        drop_last=training_args.dataloader_drop_last,
-                        num_workers=training_args.dataloader_num_workers,
-                        pin_memory=training_args.dataloader_pin_memory,
-                    )
-                except Exception as e:
-                    log_info(f"Error creating DataLoader: {str(e)}")
-                    # Try again with minimal settings
-                    log_info("Attempting to create DataLoader with minimal settings")
-                    return torch.utils.data.DataLoader(
-                        dataset,
-                        batch_size=1,  # Minimal batch size
-                        sampler=sequential_sampler,
-                        collate_fn=data_collator,
-                        num_workers=0,  # No parallel workers
-                        pin_memory=False,
-                    )
-            
-            # Override the get_train_dataloader method
-            trainer.get_train_dataloader = custom_get_train_dataloader
-            
-            # Start training
-            log_info("=== Starting Training ===")
-            try:
-                # Empty cache again right before training
-                if CUDA_AVAILABLE:
-                    torch.cuda.empty_cache()
-                    log_info("Cleared CUDA cache before training")
-                
-                # Display compact training info
-                total_steps = int((len(dataset) / (per_device_batch_size * NUM_GPUS * gradient_accumulation_steps)) * training_args.num_train_epochs)
-                log_info(f"Training plan: {len(dataset)} examples over {training_args.num_train_epochs} epochs ≈ {total_steps} steps")
-                
-                trainer.train()
-                log_info("Training completed successfully!")
-                
-                # Save the final model
-                log_info("Saving final model...")
-                trainer.save_model()
-                log_info(f"Model saved to {training_args.output_dir}")
-                
-                # Push to hub if enabled
-                if transformers_config.get("huggingface_hub", {}).get("push_to_hub", False):
-                    hub_id = transformers_config.get("huggingface_hub", {}).get("hub_model_id", "model")
-                    log_info(f"Pushing model to Hugging Face Hub as {hub_id}...")
-                    trainer.push_to_hub()
-                    log_info("Model successfully pushed to Hub")
-                
-                # Update the Hugging Face Space with current code
-                if os.environ.get("HF_TOKEN") and os.environ.get("HF_USERNAME") and os.environ.get("HF_SPACE_NAME"):
-                    update_huggingface_space()
-                    
-                return 0
-            except Exception as e:
-                logger.error(f"Training failed with error: {str(e)}")
-                # Log CUDA memory info if available in compact format
-                if CUDA_AVAILABLE:
-                    memory_info = []
-                    for i in range(NUM_GPUS):
-                        allocated = torch.cuda.memory_allocated(i) / 1024**2
-                        reserved = torch.cuda.memory_reserved(i) / 1024**2
-                        max_mem = torch.cuda.max_memory_allocated(i) / 1024**2
-                        memory_info.append(f"GPU {i}: {allocated:.1f}MB/{reserved:.1f}MB (max: {max_mem:.1f}MB)")
-                    logger.error(f"GPU memory at failure: {', '.join(memory_info)}")
-                raise
+            configure_custom_dataloader(trainer, dataset, transformers_config, training_args)
+        except Exception as e:
+            logger.error(f"Error configuring custom dataloader: {str(e)}")
+            return 1
         
+        # Run training process
+        try:
+            return run_training(trainer, tokenizer, training_args)
         except Exception as e:
-            logger.error(f"Error in main training loop: {str(e)}")
+            logger.error(f"Training failed with error: {str(e)}")
+            # Log GPU memory for debugging
+            log_gpu_memory_usage(label="Error")
+            # Print full stack trace
+            traceback.print_exc()
             return 1
     
     except Exception as e:
         logger.error(f"Error in main function: {str(e)}")
+        traceback.print_exc()
         return 1
 
 if __name__ == "__main__":