Fix the training issue and add the CPU, GPU and MPS for training.

Author: codewithdark-git

quantllm/config/model_config.py

@@ -10,7 +10,7 @@ class ModelConfig:
     model_name: str
     model_type: str = "auto"
     revision: str = "main"
-    trust_remote_code: bool = False
+    trust_remote_code: bool = True
 
     # Model architecture
     hidden_size: Optional[int] = None
@@ -42,6 +42,12 @@ class ModelConfig:
     lora_config: Optional[Dict[str, Any]] = None
     use_lora: bool = False
 
+    # CPU optimization
+    cpu_offload: bool = False
+    gradient_checkpointing: bool = False
+    bf16: bool = False  # bfloat16 support for more efficient training
+    max_memory: Optional[dict] = None  # For device specific memory limits
+
     kwargs: Optional[Dict[str, Any]] = None
     device_map: Optional[Dict[str, str]] = 'auto'  # 'auto' or specific device mapping
 
@@ -58,6 +64,20 @@ def __post_init__(self):
 
         if self.kwargs is None:
             self.kwargs = {}
+        
+        if self.load_in_4bit and self.load_in_8bit:
+            raise ValueError("Cannot use both 4-bit and 8-bit quantization simultaneously")
+        
+        # Set reasonable defaults for memory management
+        if self.max_memory is None:
+            import torch
+            if torch.cuda.is_available():
+                # Leave some GPU memory free for system
+                total_memory = torch.cuda.get_device_properties(0).total_memory
+                self.max_memory = {0: f"{int(total_memory * 0.85 / 1024**3)}GiB"}
+            else:
+                # Default CPU memory limit
+                self.max_memory = {"cpu": "16GiB"}
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert configuration to dictionary."""
@@ -88,6 +108,10 @@ def to_dict(self) -> Dict[str, Any]:
             "bnb_4bit_use_double_quant": self.bnb_4bit_use_double_quant,
             "lora_config": self.lora_config,
             "use_lora": self.use_lora,
+            "cpu_offload": self.cpu_offload,
+            "gradient_checkpointing": self.gradient_checkpointing,
+            "bf16": self.bf16,
+            "max_memory": self.max_memory,
             "kwargs": self.kwargs,
             "device_map": self.device_map
         }

quantllm/data/dataloader.py

@@ -4,46 +4,11 @@
 from datasets import Dataset as HFDataset
 from .dataset_preprocessor import DatasetPreprocessor
 
-class DataLoader(TorchDataLoader):
+class DataLoader:
     """
-    Custom DataLoader class for QuantLLM that inherits from torch.utils.data.DataLoader.
-    Provides additional functionality and easier integration with the QuantLLM package.
+    Custom DataLoader class for QuantLLM that wraps torch.utils.data.DataLoader.
     """
 
-    def __init__(
-        self,
-        dataset: Dataset,
-        batch_size: int = 4,
-        shuffle: bool = True,
-        num_workers: int = 4,
-        pin_memory: bool = True,
-        drop_last: bool = False,
-        **kwargs
-    ):
-        """
-        Initialize the QuantLLM DataLoader.
-        
-        Args:
-            dataset (Dataset): The dataset to load
-            batch_size (int): Number of samples per batch
-            shuffle (bool): Whether to shuffle the data
-            num_workers (int): Number of worker processes for data loading
-            pin_memory (bool): Whether to pin memory for faster data transfer to GPU
-            drop_last (bool): Whether to drop the last incomplete batch
-            **kwargs: Additional arguments to pass to the DataLoader
-        """
-        self.loader = TorchDataLoader(
-            dataset=dataset,
-            batch_size=batch_size,
-            shuffle=shuffle,
-            num_workers=num_workers,
-            pin_memory=pin_memory,
-            drop_last=drop_last,
-            **kwargs
-        )
-        self.dataset = dataset
-        self.batch_size = batch_size
-
     @staticmethod
     def validate_dataset(dataset, name: str):
         """Validate dataset."""
@@ -74,24 +39,47 @@ def from_datasets(
             if batch_size <= 0:
                 raise ValueError(f"batch_size must be positive, got {batch_size}")
 
-            # Convert HuggingFace Dataset to PyTorch Dataset if needed
-            def convert_to_torch_dataset(hf_dataset):
-                if hf_dataset is None:
+            def prepare_dataset(dataset):
+                if dataset is None:
                     return None
-                if isinstance(hf_dataset, HFDataset):
-                    return hf_dataset.with_format("torch")
-                return hf_dataset
+                    
+                if isinstance(dataset, HFDataset):
+                    # Ensure all required features are present
+                    required_features = ['input_ids', 'attention_mask', 'labels']
+                    if not all(feature in dataset.features for feature in required_features):
+                        raise ValueError(f"Dataset must contain all required features: {required_features}")
+                    
+                    # Get feature dimensions
+                    sample_len = len(dataset[0]['input_ids'])
+                    total_samples = len(dataset)
+                    
+                    # Pre-allocate tensors
+                    input_ids = torch.zeros((total_samples, sample_len), dtype=torch.long)
+                    attention_mask = torch.zeros((total_samples, sample_len), dtype=torch.long)
+                    labels = torch.zeros((total_samples, sample_len), dtype=torch.long)
+                    
+                    # Fill tensors
+                    for i in range(total_samples):
+                        input_ids[i] = torch.tensor(dataset[i]['input_ids'])
+                        attention_mask[i] = torch.tensor(dataset[i]['attention_mask'])
+                        labels[i] = torch.tensor(dataset[i]['labels'])
+                    
+                    return TensorDataset(input_ids, attention_mask, labels)
+                
+                return dataset
 
-            train_dataset = convert_to_torch_dataset(train_dataset)
-            val_dataset = convert_to_torch_dataset(val_dataset)
-            test_dataset = convert_to_torch_dataset(test_dataset)
+            train_dataset = prepare_dataset(train_dataset)
+            val_dataset = prepare_dataset(val_dataset)
+            test_dataset = prepare_dataset(test_dataset)
 
+            # Create DataLoaders with consistent batch sizes
             train_loader = TorchDataLoader(
                 train_dataset,
                 batch_size=batch_size,
                 shuffle=shuffle,
                 num_workers=num_workers,
                 pin_memory=pin_memory and torch.cuda.is_available(),
+                drop_last=True,  # Drop last incomplete batch
                 **kwargs
             ) if train_dataset is not None else None
 
@@ -101,6 +89,7 @@ def convert_to_torch_dataset(hf_dataset):
                 shuffle=False,
                 num_workers=num_workers,
                 pin_memory=pin_memory and torch.cuda.is_available(),
+                drop_last=True,  # Drop last incomplete batch
                 **kwargs
             ) if val_dataset is not None else None
 
@@ -110,83 +99,12 @@ def convert_to_torch_dataset(hf_dataset):
                 shuffle=False,
                 num_workers=num_workers,
                 pin_memory=pin_memory and torch.cuda.is_available(),
+                drop_last=True,  # Drop last incomplete batch
                 **kwargs
             ) if test_dataset is not None else None
 
             return train_loader, val_loader, test_loader
 
         except Exception as e:
             print(f"Error creating data loaders: {str(e)}")
-            raise
-            
-    @classmethod
-    def from_tensors(
-        cls,
-        input_ids,
-        attention_mask,
-        labels=None,
-        batch_size: int = 8,
-        **kwargs
-    ):
-        """Create DataLoader from tensor inputs."""
-        try:
-            if not isinstance(input_ids, torch.Tensor):
-                input_ids = torch.tensor(input_ids)
-            if not isinstance(attention_mask, torch.Tensor):
-                attention_mask = torch.tensor(attention_mask)
-                
-            if labels is not None:
-                if not isinstance(labels, torch.Tensor):
-                    labels = torch.tensor(labels)
-                dataset = TensorDataset(input_ids, attention_mask, labels)
-            else:
-                dataset = TensorDataset(input_ids, attention_mask)
-                
-            return TorchDataLoader(
-                dataset,
-                batch_size=batch_size,
-                **kwargs
-            )
-            
-        except Exception as e:
-            raise RuntimeError(f"Error creating data loader from tensors: {str(e)}")
-    
-    def get_batch(self) -> Dict[str, torch.Tensor]:
-        """
-        Get a single batch from the DataLoader.
-        
-        Returns:
-            Dict[str, torch.Tensor]: Dictionary containing the batch data
-        """
-        try:
-            batch = next(iter(self.loader))
-            return batch
-        except StopIteration:
-            raise RuntimeError("No more batches available in the DataLoader")
-            
-    def get_batch_size(self) -> int:
-        """
-        Get the current batch size of the DataLoader.
-        
-        Returns:
-            int: Current batch size
-        """
-        return self.batch_size
-        
-    def get_dataset_size(self) -> int:
-        """
-        Get the size of the underlying dataset.
-        
-        Returns:
-            int: Size of the dataset
-        """
-        return len(self.dataset)
-        
-    def get_num_batches(self) -> int:
-        """
-        Get the total number of batches in the DataLoader.
-        
-        Returns:
-            int: Total number of batches
-        """
-        return len(self.loader)
+            raise

quantllm/data/dataset_preprocessor.py

@@ -11,8 +11,8 @@ def __init__(self, tokenizer, logger=None):
         # Set pad token if not set
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
-            self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-            print("Added [PAD] token to tokenizer")
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+            print("Set pad token to eos token")
 
     def validate_datasets(self, datasets):
         """Validate input datasets."""
@@ -43,68 +43,95 @@ def tokenize_dataset(
             self.validate_datasets([train_dataset, val_dataset, test_dataset])
 
             def process_and_tokenize_batch(examples):
+                # Get texts and preprocess
                 texts = examples[text_column]
                 if not isinstance(texts, list):
                     texts = [texts]
-                
-                # Preprocess texts
                 texts = [self.preprocess_text(text) for text in texts]
 
                 try:
+                    # Tokenize with padding and truncation
+                    # Use max_length + 1 to account for the shift we'll do later
                     tokenized = self.tokenizer(
                         texts,
-                        padding=True,
+                        padding="max_length",
                         truncation=True,
-                        max_length=max_length,
-                        return_tensors="pt"
+                        max_length=max_length + 1,  # Add 1 to account for shift
+                        return_tensors=None
                     )
 
+                    input_ids = tokenized["input_ids"]
+                    attention_mask = tokenized["attention_mask"]
+                    
+                    # Now shift to create inputs and labels
+                    # inputs will be [:-1] and labels will be [1:]
+                    labels = [ids[1:] for ids in input_ids]
+                    input_ids = [ids[:-1] for ids in input_ids]
+                    attention_mask = [mask[:-1] for mask in attention_mask]
+                    
+                    # Verify all sequences have the expected length
+                    expected_length = max_length
+                    if not all(len(seq) == expected_length for seq in input_ids):
+                        raise ValueError(f"Input sequence lengths don't match. Expected {expected_length}")
+                    if not all(len(seq) == expected_length for seq in attention_mask):
+                        raise ValueError(f"Attention mask lengths don't match. Expected {expected_length}")
+                    if not all(len(seq) == expected_length for seq in labels):
+                        raise ValueError(f"Label sequence lengths don't match. Expected {expected_length}")
+                    
                     result = {
-                        "input_ids": tokenized["input_ids"],
-                        "attention_mask": tokenized["attention_mask"]
+                        "input_ids": input_ids,
+                        "attention_mask": attention_mask,
+                        "labels": labels
                     }
 
-                    if label_column and label_column in examples:
-                        result["labels"] = examples[label_column]
-                    
-                    print(f"Tokenized batch of {len(texts)} texts")  # User feedback
+                    self.logger.log_info(f"Tokenized batch of {len(texts)} texts")
                     return result
 
                 except Exception as e:
-                    print(f"Error tokenizing batch: {str(e)}")  # User feedback
+                    self.logger.log_error(f"Error tokenizing batch: {str(e)}")
                     raise
 
             # Process datasets
             train_tokenized = train_dataset.map(
                 process_and_tokenize_batch,
                 batched=True,
                 batch_size=batch_size,
-                remove_columns=train_dataset.column_names
+                remove_columns=train_dataset.column_names,
+                desc="Tokenizing training set"
             )
-            print(f"Tokenized training dataset: {len(train_tokenized)} examples")  # User feedback
+            self.logger.log_info(f"Tokenized training dataset: {len(train_tokenized)} examples")
 
             val_tokenized = None
             if val_dataset is not None:
                 val_tokenized = val_dataset.map(
                     process_and_tokenize_batch,
                     batched=True,
                     batch_size=batch_size,
-                    remove_columns=val_dataset.column_names
+                    remove_columns=val_dataset.column_names,
+                    desc="Tokenizing validation set"
                 )
-                print(f"Tokenized validation dataset: {len(val_tokenized)} examples")  # User feedback
+                self.logger.log_info(f"Tokenized validation dataset: {len(val_tokenized)} examples")
 
             test_tokenized = None
             if test_dataset is not None:
                 test_tokenized = test_dataset.map(
                     process_and_tokenize_batch,
                     batched=True,
                     batch_size=batch_size,
-                    remove_columns=test_dataset.column_names
+                    remove_columns=test_dataset.column_names,
+                    desc="Tokenizing test set"
                 )
-                print(f"Tokenized test dataset: {len(test_tokenized)} examples")  # User feedback
+                self.logger.log_info(f"Tokenized test dataset: {len(test_tokenized)} examples")
+            
+            # Set format to PyTorch tensors
+            train_tokenized.set_format("torch")
+            if val_tokenized:
+                val_tokenized.set_format("torch")
+            if test_tokenized:
+                test_tokenized.set_format("torch")
 
             return train_tokenized, val_tokenized, test_tokenized
 
         except Exception as e:
-            print(f"Error in tokenization: {str(e)}")  # User feedback
+            self.logger.log_error(f"Error in tokenization: {str(e)}")
             raise