Upload folder using huggingface_hub

run_transformers_training.py

@@ -151,13 +151,55 @@ def load_model_and_tokenizer(config):
             use_flash_attention = False
             logger.warning("Flash attention not available, falling back to standard attention")
             
-        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,  # Let Unsloth choose optimal dtype
-            device_map="auto",
-            # Don't explicitly use flash attention config here, let Unsloth handle it
-        )
+        # First detect if we have a GPU
+        if torch.cuda.is_available():
+            logger.info(f"CUDA available, found {torch.cuda.device_count()} GPU(s)")
+            device_map = "auto"
+        else:
+            logger.warning("No CUDA available, falling back to CPU")
+            device_map = {"": "cpu"}  # Force CPU placement
+        
+        # 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)")
+        else:
+            # CPU, use default dtype
+            dtype = None
+            logger.info("Using default precision (CPU)")
+        
+        # Load model with proper error handling for out-of-memory
+        try:
+            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,
+                # Don't explicitly use flash attention config here, let Unsloth handle it
+            )
+        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", {})
@@ -332,14 +374,16 @@ class SimpleDataCollator:
         self.dataset_config = dataset_config
         self.stats = {"processed": 0, "skipped": 0, "total_tokens": 0}
         self.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
-        self.prompt_counter = 0
         self.paper_counters = {}
         self.max_seq_length = dataset_config.get("dataset", {}).get("processing", {}).get("max_seq_length", 2048)
-        self.include_metadata = dataset_config.get("data_formatting", {}).get("metadata_handling", {}).get("include_paper_id", True)
-        self.include_chunk = dataset_config.get("data_formatting", {}).get("metadata_handling", {}).get("include_chunk_number", True)
-        self.metadata_format = dataset_config.get("data_formatting", {}).get("metadata_handling", {}).get("metadata_format", "Paper ID: {paper_id} | Chunk: {chunk_number}")
+        self.include_metadata = False  # Disable automatic metadata inclusion as it's already in content
         self.roles = dataset_config.get("data_formatting", {}).get("roles", {})
         logger.info(f"SimpleDataCollator initialized - using phi-4 chat format with max_seq_length={self.max_seq_length}")
+        logger.info("Metadata handling disabled - using metadata from content field")
+        
+        # Check if we're on GPU
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"SimpleDataCollator using device: {self.device}")
     
     def normalize_conversation(self, conversation):
         """Normalize conversation format to ensure consistent structure."""
@@ -353,6 +397,23 @@ class SimpleDataCollator:
             else:
                 return []
         
+        # Get introductory message if present (should be first and without chunk number)
+        intro_msg = None
+        for i, turn in enumerate(conversation):
+            if isinstance(turn, dict) and turn.get('content') and "[RESEARCH INTRODUCTION]" in turn.get('content', ''):
+                intro_msg = turn
+                break
+        
+        # Process introduction message first if found
+        if intro_msg:
+            normalized.append({
+                "role": "system",
+                "content": intro_msg.get('content', '')
+            })
+            # Remove intro from further processing
+            conversation = [t for t in conversation if t != intro_msg]
+        
+        # Process remaining messages
         for turn in conversation:
             # Skip empty or None entries
             if not turn:
@@ -406,23 +467,6 @@ class SimpleDataCollator:
                     self.stats["skipped"] += 1
                     continue
                 
-                # Track paper chunks
-                if paper_id not in self.paper_counters:
-                    self.paper_counters[paper_id] = 0
-                self.paper_counters[paper_id] += 1
-                
-                # Add metadata if configured
-                if self.include_metadata:
-                    # Format metadata according to configured format
-                    metadata_content = self.metadata_format.format(
-                        paper_id=paper_id, 
-                        chunk_number=self.paper_counters[paper_id]
-                    )
-                    
-                    # Add as system message if not already in conversation
-                    if not any(msg.get("role") == "system" for msg in conversation):
-                        conversation = [{"role": "system", "content": metadata_content}] + conversation
-                
                 # Format conversation with research introduction and chunk info
                 formatted_content = format_phi_chat(conversation, self.dataset_config)
                 
@@ -433,6 +477,7 @@ class SimpleDataCollator:
                     truncation=True,
                     max_length=self.max_seq_length,
                     return_tensors=None,
+                    padding=False,  # Don't pad here, we'll pad the batch later
                 )
                 
                 if len(inputs["input_ids"]) > 0:
@@ -450,7 +495,7 @@ class SimpleDataCollator:
                     log_samples = self.dataset_config.get("validation", {}).get("log_samples", 3)
                     if self.stats["processed"] <= log_samples:
                         logger.info(f"Example {self.stats['processed']} format:")
-                        logger.info(f"Paper ID: {paper_id} | Chunk: {self.paper_counters[paper_id]}")
+                        logger.info(f"Paper ID: {paper_id}")
                         logger.info(f"Token count: {len(inputs['input_ids'])}")
                         logger.info(f"Content preview:\n{formatted_content[:500]}...")
                         logger.info(f"Conversation structure: {conversation[:2]}...")
@@ -464,6 +509,7 @@ class SimpleDataCollator:
         
         if not batch["input_ids"]:
             logger.warning("Empty batch, returning dummy tensors")
+            # Return tensors on the right device
             return {
                 "input_ids": torch.zeros((1, 1), dtype=torch.long),
                 "attention_mask": torch.zeros((1, 1), dtype=torch.long),
@@ -480,8 +526,8 @@ class SimpleDataCollator:
                 batch["attention_mask"][i].extend([0] * padding_length)
                 batch["labels"][i].extend([-100] * padding_length)
         
-        # Convert to tensors
-        batch = {k: torch.tensor(v) for k, v in batch.items()}
+        # Convert to tensors on CPU first
+        batch = {k: torch.tensor(v, dtype=torch.long) for k, v in batch.items()}
         
         # Log stats periodically
         log_interval = self.dataset_config.get("validation", {}).get("log_interval", 100)
@@ -601,6 +647,18 @@ def main():
     set_seed(seed)
     logger.info(f"Set random seed to {seed}")
     
+    # Check CUDA and set environment variables for better memory management
+    if torch.cuda.is_available():
+        # Empty CUDA cache
+        torch.cuda.empty_cache()
+        # Set memory management env vars (optional)
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
+        # Log memory information
+        for i in range(torch.cuda.device_count()):
+            logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}")
+            logger.info(f"Memory Allocated: {torch.cuda.memory_allocated(i) / 1024**2:.2f} MB")
+            logger.info(f"Memory Reserved: {torch.cuda.memory_reserved(i) / 1024**2:.2f} MB")
+    
     try:
         model, tokenizer = load_model_and_tokenizer(model_config)
         logger.info("Model and tokenizer loaded successfully")
@@ -612,7 +670,7 @@ def main():
         except Exception as e:
             logger.error(f"Error loading dataset: {e}")
             return 1
-        
+            
         # Create data collator
         data_collator = SimpleDataCollator(tokenizer, dataset_config)
         
@@ -627,6 +685,13 @@ def main():
                 if (state.global_step % 50 == 0) or (current_time - self.last_log_time > 300):
                     logger.info(f"Step {state.global_step}: Loss {state.log_history[-1]['loss'] if state.log_history else 'N/A'}")
                     self.last_log_time = current_time
+                    
+            def on_train_begin(self, args, state, control, **kwargs):
+                logger.info("Training is starting...")
+                # Log memory information
+                if torch.cuda.is_available():
+                    for i in range(torch.cuda.device_count()):
+                        logger.info(f"GPU {i} Memory Allocated: {torch.cuda.memory_allocated(i) / 1024**2:.2f} MB")
         
         # Verify precision settings - ensure only one of bf16/fp16 is set, with bf16 taking precedence
         use_bf16 = model_config.get("bf16", False) or model_config.get("torch_dtype", "") == "bfloat16"
@@ -658,95 +723,73 @@ def main():
             report_to="tensorboard",
             remove_unused_columns=False,  # Keep all columns
             gradient_checkpointing=model_config.get("training", {}).get("gradient_checkpointing", True),
-            dataloader_pin_memory=False,  # Reduce memory usage
+            dataloader_pin_memory=True,  # Keep data in pinned memory for faster transfer
             optim=model_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=4,  # Sequential data loading 
+            no_cuda=False if torch.cuda.is_available() else True,  # Use CUDA if available
         )
         
-        # Create a sequential sampler to ensure dataset is processed in order
-        logger.info("Creating sequential sampler to maintain dataset order")
-        
-        # Create trainer with callback
-        logger.info("Creating trainer")
-        
-        # Check if we should resume from checkpoint
-        resume_from_checkpoint = False
-        output_dir = model_config.get("output_dir", "./results")
-        if os.path.exists(output_dir):
-            checkpoints = [folder for folder in os.listdir(output_dir) if folder.startswith("checkpoint-")]
-            if checkpoints:
-                latest_checkpoint = max(checkpoints, key=lambda x: int(x.split("-")[1]))
-                resume_from_checkpoint = os.path.join(output_dir, latest_checkpoint)
-                logger.info(f"Found checkpoint: {resume_from_checkpoint}. Training will resume from this point.")
+        # Custom dataloader to ensure no shuffling of dataset
+        # This preserves the order of chunks in papers
+        def get_train_dataloader_no_shuffle():
+            logger.info("Creating data loader with sequential sampler to maintain paper order")
+            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 * torch.cuda.device_count(), 1)
+                
+            # Use sequential sampler to preserve order
+            sequential_sampler = torch.utils.data.SequentialSampler(dataset["train"])
+            logger.info(f"Using sequential sampler for batch size {batch_size}")
+            
+            return torch.utils.data.DataLoader(
+                dataset["train"],
+                batch_size=batch_size,
+                sampler=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 trainer with custom dataloader
+        logger.info("Initializing Trainer")
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=dataset,
+            get_train_dataloader=get_train_dataloader_no_shuffle,
+            tokenizer=tokenizer,
             data_collator=data_collator,
             callbacks=[LoggingCallback()]
         )
         
-        # Override the default data loader to disable shuffling
-        # This is necessary because TrainingArguments doesn't have a direct shuffle parameter
-        def get_train_dataloader_no_shuffle():
-            """Create a train DataLoader with shuffling disabled."""
-            logger.info("Creating train dataloader with sequential sampler (no shuffling)")
-            
-            # Create a sequential sampler to ensure dataset is processed in order
-            train_sampler = torch.utils.data.SequentialSampler(dataset)
-            
-            return torch.utils.data.DataLoader(
-                dataset,
-                batch_size=training_args.per_device_train_batch_size,
-                sampler=train_sampler,  # Use sequential sampler instead of shuffle parameter
-                collate_fn=data_collator,
-                drop_last=False,
-                num_workers=0,
-                pin_memory=False
-            )
-        
-        # Replace the default data loader with our non-shuffling version
-        trainer.get_train_dataloader = get_train_dataloader_no_shuffle
-        
         # Start training
-        logger.info("Starting training")
-        logger.info(f"Processing with batch size = {training_args.per_device_train_batch_size}, each entry processed independently")
-        
-        # Create a lock file to indicate training is in progress
-        lock_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "TRAINING_IN_PROGRESS.lock")
-        with open(lock_file, "w") as f:
-            f.write(f"Training started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
-            f.write(f"Expected completion: After {training_args.num_train_epochs} epochs\n")
-            f.write("DO NOT UPDATE OR RESTART THIS SPACE UNTIL TRAINING COMPLETES\n")
-        logger.info(f"Created lock file: {lock_file}")
-        
+        logger.info("Starting training process")
         try:
-            trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+            trainer.train()
             logger.info("Training completed successfully")
             
-            # Save model
-            if model_config.get("push_to_hub", False):
-                logger.info(f"Pushing model to hub: {model_config.get('hub_model_id')}")
+            # Save the final model
+            logger.info("Saving final model")
+            trainer.save_model()
+            
+            # Push to hub if enabled
+            if model_config.get("huggingface_hub", {}).get("push_to_hub", False):
+                logger.info("Pushing model to Hugging Face Hub")
                 trainer.push_to_hub()
-                logger.info("Model pushed to hub successfully")
-            else:
-                logger.info(f"Saving model to {model_config.get('output_dir', './results')}")
-                trainer.save_model()
-                logger.info("Model saved successfully")
+                
+            return 0
         except Exception as e:
             logger.error(f"Training failed with error: {str(e)}")
+            # Log CUDA memory info if available
+            if torch.cuda.is_available():
+                for i in range(torch.cuda.device_count()):
+                    logger.info(f"GPU {i} Memory Allocated: {torch.cuda.memory_allocated(i) / 1024**2:.2f} MB")
+                    logger.info(f"GPU {i} Memory Reserved: {torch.cuda.memory_reserved(i) / 1024**2:.2f} MB")
             raise
-        finally:
-            # Remove the lock file when training completes or fails
-            if os.path.exists(lock_file):
-                os.remove(lock_file)
-                logger.info(f"Removed lock file: {lock_file}")
-            
-            return 0
-    
+        
     except Exception as e:
         logger.error(f"Error in main training loop: {str(e)}")
         return 1