Merge branch 'quic:main' into main

Author: vbaddi

QEfficient/cloud/finetune.py

@@ -103,6 +103,18 @@ def main(**kwargs):
     # print the datatype of the model parameters
     # print(get_parameter_dtypes(model))
 
+    # Note: Need to call this before calling PeftModel.from_pretrained or get_peft_model.
+    # Because, both makes model.is_gradient_checkpointing = True which is used in peft library to
+    # apply gradient checkpointing related hooks to the input embeddings. Without this we will get
+    # "No inf checks were recorded for this optimizer." error.
+    # Enable gradient checkpointing
+    if train_config.gradient_checkpointing:
+        # Note: below attribute and method is only available in HuggingFace Transformer models.
+        if hasattr(model, "supports_gradient_checkpointing") and model.supports_gradient_checkpointing:
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"preserve_rng_state": False})
+        else:
+            raise RuntimeError("Given model doesn't support gradient checkpointing. Please disable it and run it.")
+
     if train_config.use_peft:
         # Load the pre-trained peft model checkpoint and setup its configuration
         if train_config.from_peft_checkpoint:

QEfficient/finetune/configs/training.py

@@ -15,6 +15,7 @@ class train_config:
     batch_size_training: int = 1
     context_length: int = None
     gradient_accumulation_steps: int = 4
+    gradient_checkpointing: bool = False
     num_epochs: int = 1
     max_train_step: int = 0
     max_eval_step: int = 0

QEfficient/finetune/dataset/dataset_config.py

@@ -21,6 +21,9 @@
 from QEfficient.finetune.dataset.samsum_dataset import (
     get_preprocessed_samsum as get_samsum_dataset,
 )
+from QEfficient.finetune.dataset.samsum_dataset import (
+    get_samsum_collate_fn,
+)
 
 DATASET_PREPROC = {
     "alpaca_dataset": partial(get_alpaca_dataset),
@@ -29,4 +32,7 @@
     "gsm8k_dataset": get_gsm8k_dataset,
     "custom_dataset": get_custom_dataset,
 }
-DATALOADER_COLLATE_FUNC = {"custom_dataset": get_data_collator}
+DATALOADER_COLLATE_FUNC = {
+    "custom_dataset": get_data_collator,
+    "samsum_dataset": get_samsum_collate_fn,
+}

QEfficient/finetune/dataset/samsum_dataset.py

@@ -6,6 +6,8 @@
 # -----------------------------------------------------------------------------
 
 import datasets
+import torch
+from torch.nn.utils.rnn import pad_sequence
 
 
 def get_preprocessed_samsum(dataset_config, tokenizer, split, context_length=None):
@@ -46,3 +48,22 @@ def tokenize_add_label(sample):
     dataset = dataset.map(tokenize_add_label, remove_columns=list(dataset.features))
 
     return dataset
+
+
+def collate_fn(batch):
+    eos_token = batch[0]["input_ids"][-1]
+
+    input_ids = pad_sequence(
+        [torch.tensor(b["input_ids"], dtype=torch.int32) for b in batch], batch_first=True, padding_value=eos_token
+    )
+    attn_mask = pad_sequence(
+        [torch.tensor(b["attention_mask"], dtype=torch.int32) for b in batch], batch_first=True, padding_value=0
+    )
+    labels = pad_sequence(
+        [torch.tensor(b["labels"], dtype=torch.long) for b in batch], batch_first=True, padding_value=eos_token
+    )
+    return {"input_ids": input_ids, "attention_mask": attn_mask, "labels": labels}
+
+
+def get_samsum_collate_fn(dataset_processer, dataset_config):
+    return collate_fn

QEfficient/finetune/utils/train_utils.py

@@ -83,6 +83,7 @@ def train(
     best_val_loss = float("inf")
     total_train_steps = 0
     max_steps_reached = False  # Flag to indicate max training steps reached
+    device_type = device.split(":")[0]
 
     tensorboard_updates = None
     if train_config.enable_ddp:
@@ -95,7 +96,7 @@ def train(
         if device.startswith("qaic"):
             scaler = QAicGradScaler()
         else:
-            scaler = GradScaler()
+            scaler = GradScaler(device_type)
 
     loss_0_counter = torch.tensor([0]).to(device)
 
@@ -177,10 +178,7 @@ def train(
                         # adjust atol & rtol this as required
                         atol=1e-1,
                         use_ref_output_on_mismatch=True,
-                        # report all mismatches
-                        max_failures=None,
-                        # generate unittest for each op once
-                        repeat_same_op=True,
+                        filter_config=qaic_debug.DispatchFilterConfig.default(device),
                         dump_root_dir=train_config.dump_root_dir + str(step),
                     ) as verifier:
                         loss = model(**batch).loss  # Forward call
@@ -296,8 +294,6 @@ def train(
                 eval_ppl, eval_epoch_loss, temp_val_loss, temp_step_perplexity = evaluation(
                     model, train_config, eval_dataloader, local_rank, tokenizer, device
                 )
-                dist.barrier()
-                dist.all_reduce(eval_epoch_loss, op=dist.ReduceOp.SUM)
                 if local_rank == 0:
                     tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
 

QEfficient/transformers/cache_utils.py

@@ -9,7 +9,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 import torch
-from transformers.cache_utils import DynamicCache
+from transformers.cache_utils import DynamicCache, EncoderDecoderCache
 
 from QEfficient.customop import (
     CtxGatherFunc,
@@ -181,3 +181,28 @@ def update3D(
             v_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out)
 
         return k_out, v_out
+
+
+class QEffEncoderDecoderCache(EncoderDecoderCache):
+    """
+    Updated the `EncoderDecoderCache` to use the `QEffDynamicCache` for both self-attention and cross-attention caches.
+    """
+
+    @classmethod
+    def from_legacy_cache(
+        cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    ) -> "EncoderDecoderCache":
+        """Converts a cache in the legacy cache format into an equivalent `EncoderDecoderCache`."""
+        cache = cls(
+            self_attention_cache=QEffDynamicCache(),
+            cross_attention_cache=QEffDynamicCache(),
+        )
+        if past_key_values is not None:
+            for layer_idx in range(len(past_key_values)):
+                key_states, value_states = past_key_values[layer_idx][:2]
+                cache.self_attention_cache.update(key_states, value_states, layer_idx)
+                if len(past_key_values[layer_idx]) > 2:
+                    key_states, value_states = past_key_values[layer_idx][2:]
+                    cache.cross_attention_cache.update(key_states, value_states, layer_idx)
+                    cache.is_updated[layer_idx] = True
+        return cache

QEfficient/transformers/models/codegen/modeling_codegen.py

@@ -10,20 +10,18 @@
 from typing import Optional, Tuple, Union
 
 import torch
-import torch.utils.checkpoint
 from torch import nn
-from torch.nn import CrossEntropyLoss
-from transformers.cache_utils import Cache, DynamicCache
+from transformers.cache_utils import Cache
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.models.codegen.modeling_codegen import (
     CodeGenAttention,
     CodeGenBlock,
     CodeGenForCausalLM,
     CodeGenModel,
     apply_rotary_pos_emb,
-    logger,
 )
 
+from QEfficient.transformers.cache_utils import QEffDynamicCache
 from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask
 
 
@@ -133,7 +131,7 @@ def forward(
                 "position_ids": position_ids,
                 "batch_index": batch_index,
             }
-            pkv = DynamicCache()
+            pkv = QEffDynamicCache()
             pkv.key_cache.append(past_key_value[0])
             pkv.value_cache.append(past_key_value[1])
             key, value = pkv.update(key, value, 0, cache_kwargs)
@@ -261,14 +259,6 @@ def forward(
 
         output_shape = input_shape + (hidden_states.size(-1),)
 
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
-                    "`use_cache=False`..."
-                )
-                use_cache = False
-
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
@@ -279,41 +269,17 @@ def forward(
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
-            if self.gradient_checkpointing and self.training:
-                outputs = self._gradient_checkpointing_func(
-                    block.__call__,
-                    hidden_states,
-                    None,
-                    attention_mask,
-                    position_ids,
-                    head_mask[i],
-                    use_cache,
-                    output_attentions,
-                    cache_position,
-                )
-            elif batch_index is not None:
-                outputs = block(
-                    hidden_states=hidden_states,
-                    layer_past=layer_past,
-                    batch_index=batch_index,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    head_mask=head_mask[i],
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    cache_position=cache_position,
-                )
-            else:
-                outputs = block(
-                    hidden_states=hidden_states,
-                    layer_past=layer_past,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    head_mask=head_mask[i],
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    cache_position=cache_position,
-                )
+            outputs = block(
+                hidden_states=hidden_states,
+                layer_past=layer_past,
+                batch_index=batch_index,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                head_mask=head_mask[i],
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                cache_position=cache_position,
+            )
 
             hidden_states = outputs[0]
             if use_cache is True:
@@ -398,25 +364,8 @@ def forward(
         hidden_states = transformer_outputs[0][torch.arange(position_ids.shape[0]).view(-1, 1), logit_index]
         lm_logits = self.lm_head(hidden_states)
 
-        loss = None
-        if labels is not None:
-            # move labels to correct device to enable model parallelism
-            labels = labels.to(lm_logits.device)
-            # Shift so that tokens < n predict n
-            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-
-            loss = loss.to(hidden_states.dtype)
-
-        if not return_dict:
-            output = (lm_logits,) + transformer_outputs[1:]
-            return ((loss,) + output) if loss is not None else output
-
         return CausalLMOutputWithPast(
-            loss=loss,
+            loss=None,
             logits=lm_logits,
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,