Merge branch 'main' into pp_ddp

Signed-off-by: Mamta Singh <[email protected]>

Author: quic-mamta

QEfficient/finetune/utils/dataset_utils.py

@@ -6,6 +6,7 @@
 # -----------------------------------------------------------------------------
 
 from typing import Dict, List, Tuple
+import logging
 
 import datasets
 import torch
@@ -69,6 +70,11 @@ def get_dataloader_kwargs(train_config, dataset, dataset_processer, split):
 
 
 def padding_dataset(train_config, dataset, batch_size):
+    num_replicas = get_num_ddp_devices()
+    remainder = len(dataset) % (num_replicas * batch_size)
+    if remainder == 0:
+        return dataset
+
     if train_config.enable_ddp and train_config.enable_sorting_for_ddp:
         if isinstance(dataset, datasets.Dataset):
             # Hugging Face Dataset transformation
@@ -80,24 +86,26 @@ def padding_dataset(train_config, dataset, batch_size):
 
     dummy_row = next(iter(dataset))
     dummy_row["labels"] = torch.tensor([-100] * len(dummy_row["labels"]))
-    padding_size = 0
-    num_replicas = get_num_ddp_devices()
-    remainder = len(dataset) % (num_replicas * batch_size)
-    padding_size = (num_replicas * batch_size) - remainder
 
+    padding_size = (num_replicas * batch_size) - remainder
     dummy_data = [dummy_row.copy() for _ in range(padding_size)]
     dummy_dataset = datasets.Dataset.from_list(dummy_data)
     if isinstance(dataset, datasets.Dataset):
         combined_dataset = datasets.concatenate_datasets([dataset, dummy_dataset])
     else:
         combined_dataset = dataset + list(dummy_dataset)
+
+    logger.log_rank_zero("Padding dataset to make it divisible by batch_size * num_devices.", logging.DEBUG)
+    logger.log_rank_zero(f"Length of dataset before padding: {len(dataset)}", logging.DEBUG)
+    logger.log_rank_zero(f"Length of dataset after padding: {len(combined_dataset)}", logging.DEBUG)
     return combined_dataset
 
 
 def get_dataloader(tokenizer, dataset_config, train_config, split: str = "train"):
     dataset = get_preprocessed_dataset(tokenizer, dataset_config, split, context_length=train_config.context_length)
 
     batch_size = train_config.train_batch_size if split == "train" else train_config.val_batch_size
+
     dataset = padding_dataset(train_config, dataset, batch_size)
 
     dl_kwargs = get_dataloader_kwargs(train_config, dataset, tokenizer, split)

QEfficient/finetune/utils/train_utils.py

@@ -20,6 +20,7 @@
 from QEfficient.finetune.utils.helper import (
     Task_Mode,
     get_autocast_ctx,
+    get_num_ddp_devices,
     get_op_verifier_ctx,
     is_rank_zero,
     save_to_json,
@@ -67,8 +68,8 @@ def train(
 
     train_metric = []
     train_loss = []
-    val_metric = []
-    val_loss = []
+    eval_metric = []
+    eval_loss = []
 
     if train_config.save_metrics:
         if not os.path.exists(train_config.output_dir):
@@ -78,13 +79,13 @@ def train(
         )
         train_step_metric = []
         train_step_loss = []
-        val_step_loss = []
-        val_step_metric = []
+        eval_step_loss = []
+        eval_step_metric = []
 
     epoch_times = []
     checkpoint_times = []
     results = {}
-    best_val_loss = float("inf")
+    best_eval_loss = float("inf")
     total_train_steps = 0
     max_steps_reached = False  # Flag to indicate max training steps reached
 
@@ -132,8 +133,7 @@ def train(
                 continue
 
         logger.log_rank_zero(f"Starting epoch {epoch + 1}/{train_config.num_epochs}")
-        if train_config.max_train_step > 0:
-            logger.log_rank_zero(f"Max train steps : {train_config.max_train_step}")
+
         if max_steps_reached:
             break
 
@@ -170,6 +170,11 @@ def train(
 
             if train_config.max_train_step > 0 and total_train_steps > train_config.max_train_step:
                 max_steps_reached = True
+                logger.log_rank_zero(
+                    "Maximum training steps reached "
+                    f"(max_train_step={train_config.max_train_step}). Stopping "
+                    "the training process."
+                )
                 break
             batch = {k: v.to(device) for k, v in batch.items()}  # move the batch elements to qaic device
 
@@ -207,6 +212,7 @@ def train(
                 logger.info("Mismatches detected:", verifier.get_perop_mismatch_count())
 
             total_loss += loss.detach().float()
+
             if is_rank_zero():
                 tensorboard_updates.add_scalars("loss", {"train": loss}, total_train_steps)
                 if loss <= train_config.convergence_loss:
@@ -219,10 +225,10 @@ def train(
             if train_config.save_metrics:
                 train_step_loss.append(loss.detach().float().item())
                 if train_config.task_mode == Task_Mode.SEQ_CLASSIFICATION:
-                    step_metric_val = float(acc_helper.compute())
+                    step_metric_value = float(acc_helper.compute())
                 else:
-                    step_metric_val = float(torch.exp(loss.detach().float()))
-                train_step_metric.append(step_metric_val)
+                    step_metric_value = float(torch.exp(loss.detach().float()))
+                train_step_metric.append(step_metric_value)
 
             # Accumulate gradients
             complete_accum_steps = (
@@ -250,15 +256,17 @@ def train(
                 pbar.update(1)
 
             # Save the trained checkpoints for every given steps
-            if step % train_config.intermediate_step_save == 0:
+            if (step + 1) % train_config.intermediate_step_save == 0:
                 qaic_profile.stop_profiling(device) if train_config.use_profiler else None
                 if train_config.enable_ddp:
                     if dist.get_rank() == 0:
                         model.module.save_pretrained(
-                            train_config.output_dir + f"/trained_weights/epoch_{epoch + 1}/step_{step}"
+                            train_config.output_dir + f"/trained_weights/epoch_{epoch + 1}/step_{step + 1}"
                         )
                 else:
-                    model.save_pretrained(train_config.output_dir + f"/trained_weights/epoch_{epoch + 1}/step_{step}")
+                    model.save_pretrained(
+                        train_config.output_dir + f"/trained_weights/epoch_{epoch + 1}/step_{step + 1}"
+                    )
 
             pbar.set_description(
                 f"Training Epoch: {epoch + 1}/{train_config.num_epochs}, step {step + 1}/{len(train_dataloader)} completed (loss: {loss.detach().float()})"
@@ -270,10 +278,10 @@ def train(
                     train_loss,
                     train_step_metric,
                     train_metric,
-                    val_step_loss,
-                    val_loss,
-                    val_step_metric,
-                    val_metric,
+                    eval_step_loss,
+                    eval_loss,
+                    eval_step_metric,
+                    eval_metric,
                 )
             if loss_0_counter.item() == train_config.convergence_counter:
                 logger.log_rank_zero(
@@ -285,44 +293,64 @@ def train(
         epoch_end_time = time.perf_counter() - epoch_start_time
         epoch_times.append(epoch_end_time)
 
-        if loss_0_counter.item() == train_config.convergence_counter:
-            if train_config.use_peft and train_config.from_peft_checkpoint and epoch == intermediate_epoch:
-                train_epoch_loss = (
-                    0.0
-                    if total_loss == 0.0
-                    else total_loss / (step - intermediate_step - num_dummy_samples / train_config.train_batch_size)
-                )
-            else:
-                train_epoch_loss = (
-                    0.0
-                    if total_loss == 0.0
-                    else total_loss / (step + 1 - num_dummy_samples / train_config.train_batch_size)
-                )
+        if train_config.use_peft and train_config.from_peft_checkpoint and epoch == intermediate_epoch:
+            train_epoch_loss = (
+                0.0
+                if total_loss == 0.0
+                else total_loss / (step - intermediate_step - (num_dummy_samples / train_config.train_batch_size))
+            )
         else:
-            if train_config.use_peft and train_config.from_peft_checkpoint and epoch == intermediate_epoch:
-                train_epoch_loss = (
-                    0.0
-                    if total_loss == 0.0
-                    else total_loss / (step - intermediate_step - (num_dummy_samples / train_config.train_batch_size))
-                )
-            else:
-                train_epoch_loss = (
-                    0.0
-                    if total_loss == 0.0
-                    else total_loss / (step + 1 - (num_dummy_samples / train_config.train_batch_size))
-                )
+            train_epoch_loss = (
+                0.0
+                if total_loss == 0.0
+                else total_loss / (step + 1 - (num_dummy_samples / train_config.train_batch_size))
+            )
+
         if train_config.task_mode == Task_Mode.SEQ_CLASSIFICATION:
-            metric_val = acc_helper.compute()
+            train_epoch_metric = acc_helper.compute()
             acc_helper.reset()
         else:
-            metric_val = torch.exp(train_epoch_loss)
+            train_epoch_metric = torch.exp(train_epoch_loss)
 
-        train_metric.append(float(metric_val))
+        train_metric.append(float(train_epoch_metric))
         train_loss.append(float(train_epoch_loss))
 
+        if train_config.enable_ddp:
+            dist.all_reduce(train_epoch_loss, op=dist.ReduceOp.SUM)
+            train_epoch_loss /= get_num_ddp_devices()
+            dist.all_reduce(train_epoch_metric, op=dist.ReduceOp.SUM)
+            train_epoch_metric /= get_num_ddp_devices()
+
         # Update the learning rate as needed
         lr_scheduler.step()
 
+        if train_config.run_validation:
+            eval_epoch_loss, eval_epoch_metric, step_loss, step_metric = evaluation(
+                model, train_config, eval_dataloader, device
+            )
+
+            if is_rank_zero():
+                tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
+            if train_config.save_metrics:
+                eval_step_loss.extend(step_loss)
+                eval_step_metric.extend(step_metric)
+                eval_loss.append(float(eval_epoch_loss))
+                eval_metric.append(float(eval_epoch_metric))
+
+            if train_config.enable_ddp:
+                dist.all_reduce(eval_epoch_loss, op=dist.ReduceOp.SUM)
+                eval_epoch_loss /= get_num_ddp_devices()
+                dist.all_reduce(eval_epoch_metric, op=dist.ReduceOp.SUM)
+                eval_epoch_metric /= get_num_ddp_devices()
+
+            if eval_epoch_loss < best_eval_loss:
+                best_eval_loss = eval_epoch_loss
+                logger.log_rank_zero(f"Best eval loss on epoch {epoch + 1} is {best_eval_loss:.4f}")
+
+            logger.log_rank_zero(
+                f"Epoch {epoch + 1}: Eval Loss: {eval_epoch_loss.detach().cpu():.4f}, Eval metric: {eval_epoch_metric.detach().cpu():.4f}"
+            )
+
         # saving the adapters after completion of each epoch
         if train_config.save_model:
             if train_config.enable_ddp:
@@ -331,25 +359,10 @@ def train(
             else:
                 model.save_pretrained(train_config.output_dir + f"/complete_epoch_{epoch + 1}")
 
-        if train_config.run_validation:
-            eval_epoch_loss, eval_metric, temp_val_loss, temp_step_metric = evaluation(
-                model, train_config, eval_dataloader, device
-            )
-            if eval_epoch_loss < best_val_loss:
-                best_val_loss = eval_epoch_loss
-                logger.log_rank_zero(f"best eval loss on epoch {epoch + 1} is {best_val_loss}")
-
-            if is_rank_zero():
-                tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
-
-            if train_config.save_metrics:
-                val_step_loss.extend(temp_val_loss)
-                val_step_metric.extend(temp_step_metric)
-            val_loss.append(float(eval_epoch_loss))
-            val_metric.append(float(eval_metric))
         logger.log_rank_zero(
-            f"Epoch {epoch + 1}: train_metric={metric_val:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
+            f"Epoch {epoch + 1}: Train epoch loss: {train_epoch_loss:.4f}, Train metric: {train_epoch_metric:.4f}, Epoch time {epoch_end_time:.2f} sec"
         )
+
         # Saving the results every epoch to plot later
         if train_config.save_metrics:
             save_to_json(
@@ -358,24 +371,19 @@ def train(
                 train_loss,
                 train_step_metric,
                 train_metric,
-                val_step_loss,
-                val_loss,
-                val_step_metric,
-                val_metric,
+                eval_step_loss,
+                eval_loss,
+                eval_step_metric,
+                eval_metric,
             )
     avg_epoch_time = sum(epoch_times) / len(epoch_times)
     avg_checkpoint_time = sum(checkpoint_times) / len(checkpoint_times) if len(checkpoint_times) > 0 else 0
-    avg_train_metric = sum(train_metric) / len(train_metric)
-    avg_train_loss = sum(train_loss) / len(train_loss)
-    if train_config.run_validation:
-        avg_eval_metric = sum(val_metric) / len(val_metric)
-        avg_eval_loss = sum(val_loss) / len(val_loss)
 
-    results["avg_train_metric"] = avg_train_metric
-    results["avg_train_loss"] = avg_train_loss
+    results["last_epoch_train_loss"] = train_epoch_loss.cpu()
+    results["last_epoch_train_metric"] = train_epoch_metric.cpu()
     if train_config.run_validation:
-        results["avg_eval_metric"] = avg_eval_metric
-        results["avg_eval_loss"] = avg_eval_loss
+        results["last_epoch_eval_loss"] = eval_epoch_loss.cpu()
+        results["last_epoch_eval_metric"] = eval_epoch_metric.cpu()
     results["avg_epoch_time"] = avg_epoch_time
     results["avg_checkpoint_time"] = avg_checkpoint_time
     if train_config.save_metrics:
@@ -391,7 +399,7 @@ def evaluation(model, train_config, eval_dataloader, device):
         model: The model to evaluate
         eval_dataloader: The dataloader containing the evaluation data
 
-    Returns: eval_epoch_loss, eval_metric, eval_step_loss, eval_step_metric
+    Returns: eval_epoch_loss, eval_epoch_metric, eval_step_loss, eval_step_metric
     """
     if train_config.enable_ddp:
         dist.barrier()
@@ -408,17 +416,17 @@ def evaluation(model, train_config, eval_dataloader, device):
     # special handling for qaic device and dtype
     # model.to(device)
 
-    val_step_loss = []
-    val_step_metric = []
+    eval_step_loss = []
+    eval_step_metric = []
 
-    eval_loss = 0.0  # Initialize evaluation loss
+    eval_loss = torch.tensor(0.0, dtype=torch.float32, device=device)  # Initialize evaluation loss
     device_type = torch.device(device).type
 
     num_dummy_samples = 0
     autocast_ctx = get_autocast_ctx(train_config.use_autocast, device_type, dtype=torch.float16)
     for step, batch in enumerate(tqdm(eval_dataloader, colour="green", desc="evaluating Epoch", dynamic_ncols=True)):
         #  stop when the maximum number of eval steps is reached
-        if train_config.max_eval_step > 0 and step > train_config.max_eval_step:
+        if train_config.max_eval_step > 0 and step >= train_config.max_eval_step:
             break
         for key in batch.keys():
             batch[key] = batch[key].to(device)
@@ -445,29 +453,27 @@ def evaluation(model, train_config, eval_dataloader, device):
                 logits = outputs.logits
                 labels = batch["labels"][:, 0]
                 preds = torch.nn.functional.softmax(logits, dim=-1)
-                val_acc = acc_helper.forward(preds, labels)
-                metric_val = val_acc.detach().float().item()
+                eval_acc = acc_helper.forward(preds, labels)
+                metric_value = eval_acc.detach().float().item()
             else:
-                metric_val = float(torch.exp(loss.detach().float()))
+                metric_value = float(torch.exp(loss.detach().float()))
 
             if train_config.save_metrics:
-                val_step_loss.append(loss.detach().float().item())
-                val_step_metric.append(metric_val)
+                eval_step_loss.append(loss.detach().float().item())
+                eval_step_metric.append(metric_value)
 
             eval_loss += loss.detach().float()
+
     # Compute average loss and metric
     eval_epoch_loss = (
         0.0 if eval_loss == 0.0 else eval_loss / (step + 1 - num_dummy_samples / train_config.val_batch_size)
     )
     if train_config.task_mode == Task_Mode.SEQ_CLASSIFICATION:
-        eval_metric = acc_helper.compute()
+        eval_epoch_metric = acc_helper.compute()
     else:
-        eval_metric = torch.exp(eval_epoch_loss)
-
-    # Print evaluation metrics
-    logger.log_rank_zero(f"{eval_metric.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
+        eval_epoch_metric = torch.exp(eval_epoch_loss)
 
-    return eval_epoch_loss, eval_metric, val_step_loss, val_step_metric
+    return eval_epoch_loss, eval_epoch_metric, eval_step_loss, eval_step_metric
 
 
 def print_model_size(model) -> None: