Upload folder using huggingface_hub

run_transformers_training.py

@@ -262,13 +262,22 @@ def load_model_and_tokenizer(config):
             
         # 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=unsloth_config.get("dropout", 0.05),
+            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),
@@ -279,11 +288,19 @@ def load_model_and_tokenizer(config):
         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")
-                tokenizer.chat_template = template
-                logger.info("Set phi chat template")
+                # 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)")
+                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")
         
         # Ensure proper token settings
         if tokenizer.pad_token_id is None:
@@ -308,7 +325,23 @@ def load_dataset_with_mapping(dataset_config):
             raise ValueError("Dataset name not provided in configuration")
         
         logger.info(f"Loading pre-processed dataset {dataset_name}, split {dataset_split}")
-        dataset = load_dataset(dataset_name, split=dataset_split)
+        
+        try:
+            dataset = load_dataset(dataset_name, split=dataset_split)
+            
+            # 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)")
+                
+        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
         
         # Apply minimal processing since the dataset has already been properly structured
         # Just perform validation to ensure required fields exist
@@ -765,296 +798,311 @@ def main():
     # Set up logging
     logger.info("Starting training process")
     
-    # Check dependencies first, before any other operations
-    if not check_dependencies():
-        logger.error("Aborting due to missing critical dependencies")
-        return 1
-    
-    # Parse arguments
-    args = parse_args()
-    
-    # Load environment variables
-    load_env_variables()
-    
-    # Validate Hugging Face credentials if we're going to use them
-    validate_huggingface_credentials()
-    
-    # Load configuration
     try:
-        transformers_config = load_configs(args.config)
-        hardware_config = transformers_config.get("hardware", {})
-        dataset_config = transformers_config.get("dataset", {})
-        logger.info("Configuration loaded successfully")
-    except Exception as e:
-        logger.error(f"Error loading configuration: {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")
-    
-    # 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)
-        
-        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")
-    
-    try:
-        log_info("Loading dataset...")
-        dataset = load_dataset_with_mapping(dataset_config)
-        log_info(f"Dataset loaded with {len(dataset)} examples")
+        # Check dependencies first, before any other operations
+        if not check_dependencies():
+            logger.error("Aborting due to missing critical dependencies")
+            return 1
+        
+        # Parse arguments
+        args = parse_args()
         
-        # Minimal validation before proceeding
-        if dataset is None or len(dataset) == 0:
-            logger.error("Dataset is empty or None! Cannot proceed with training.")
+        # Load environment variables
+        load_env_variables()
+        
+        # Validate Hugging Face credentials if we're going to use them
+        validate_huggingface_credentials()
+        
+        # Load configuration
+        try:
+            transformers_config = load_configs(args.config)
+            hardware_config = transformers_config.get("hardware", {})
+            dataset_config = transformers_config.get("dataset", {})
+            logger.info("Configuration loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading configuration: {e}")
             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")
+        # 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(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")
-        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=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
-            # Only add FSDP if we're in distributed mode with FSDP strategy
-            fsdp=fsdp_config if is_distributed and multi_gpu_strategy == "fsdp" else None,
-        )
+            log_info("Running in non-distributed mode (single process)")
         
-        # Create sequential sampler to maintain original dataset order
-        sequential_sampler = torch.utils.data.SequentialSampler(dataset)
-        
-        # Initialize trainer first
-        log_info("Initializing Trainer")
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=dataset,  # We'll override this with our custom dataloader
-            data_collator=data_collator,
-            callbacks=[LoggingCallback()],
-        )
+        # 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)
         
-        # 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")
+        # Empty CUDA cache to ensure clean state
+        if CUDA_AVAILABLE:
+            torch.cuda.empty_cache()
+            log_info("Cleared CUDA cache")
+        
+        # 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)
             
-            # Create a simple sequential sampler
-            sequential_sampler = torch.utils.data.SequentialSampler(dataset)
+            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")
+        
+        try:
+            log_info("Loading dataset...")
+            dataset = load_dataset_with_mapping(dataset_config)
             
-            # Verification of sequence preservation flags - simplified
-            data_loading_config = dataset_config.get("data_loading", {})
-            shuffle_enabled = data_loading_config.get("shuffle", False)
+            # 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")
             
-            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
+            # 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
             
-            # Log our approach clearly
-            log_info("Using SequentialSampler to guarantee dataset order is preserved based on prompt_number")
+            # Create data collator
+            data_collator = SimpleDataCollator(tokenizer, dataset_config)
             
-            # Verify column order
-            expected_order = ["prompt_number", "article_id", "conversations"]
-            if hasattr(dataset, 'column_names'):
-                actual_order = dataset.column_names
-                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")
+            # 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
             
-            log_info("Dataset is pre-processed with prompt_number field indicating the correct sequence")
+            # 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'}")
             
-            # Calculate batch size based on device availability
-            if getattr(training_args, "no_cuda", False):
-                batch_size = training_args.per_device_train_batch_size
+            # 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:
-                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}")
+                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)
             
-            # Return DataLoader with sequential sampler
-            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,
+            # Set up training arguments
+            log_info("Setting up 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=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
+                # Only add FSDP if we're in distributed mode with FSDP strategy
+                fsdp=fsdp_config if is_distributed and multi_gpu_strategy == "fsdp" else None,
             )
-        
-        # 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")
+            # Create sequential sampler to maintain original dataset order
+            sequential_sampler = torch.utils.data.SequentialSampler(dataset)
             
-            trainer.train()
-            log_info("Training completed successfully!")
+            # Initialize trainer first
+            log_info("Initializing Trainer")
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=dataset,  # We'll override this with our custom dataloader
+                data_collator=data_collator,
+                callbacks=[LoggingCallback()],
+            )
             
-            # Save the final model
-            log_info("Saving final model...")
-            trainer.save_model()
-            log_info(f"Model saved to {training_args.output_dir}")
+            # 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")
+                
+                # 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
+                expected_order = ["prompt_number", "article_id", "conversations"]
+                if hasattr(dataset, 'column_names'):
+                    actual_order = dataset.column_names
+                    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")
+                
+                # 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
+                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,
+                )
             
-            # 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")
+            # Override the get_train_dataloader method
+            trainer.get_train_dataloader = custom_get_train_dataloader
             
-            # 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()
+            # 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!")
                 
-            return 0
+                # 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
+        
         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
+            logger.error(f"Error in main training loop: {str(e)}")
+            return 1
     
     except Exception as e:
-        logger.error(f"Error in main training loop: {str(e)}")
+        logger.error(f"Error in main function: {str(e)}")
         return 1
 
 if __name__ == "__main__":

transformers_config.json

@@ -46,7 +46,7 @@
     "enabled": true,
     "r": 32,
     "alpha": 16,
-    "dropout": 0.05,
+    "dropout": 0,
     "target_modules": [
       "q_proj",
       "k_proj",