From 33fd672924e82ee78b2777c64a0367378692fadf Mon Sep 17 00:00:00 2001
From: zzhhjjj <z785566960@gmail.com>
Date: Mon, 22 Apr 2024 14:39:51 +0000
Subject: [PATCH 01/33] readme

---
 examples/mamba/README.md | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/examples/mamba/README.md b/examples/mamba/README.md
index 5c31d07f..8eefa9c2 100644
--- a/examples/mamba/README.md
+++ b/examples/mamba/README.md
@@ -18,6 +18,18 @@ pip install -r requirements.txt
 
 > https://wandb.ai/bouteille/test/reports/Mamba-loss--Vmlldzo2OTgwNDM5
 
+## Bug related to nanotron
+Encountered the following issue when ran train_mamba.sh:   
+```
+causal_conv1d_cuda.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZNK3c1017SymbolicShapeMeta18init_is_contiguousEv
+```
+Solved this by doing:    
+pip uninstall mamba-ssm   
+pip install causal_conv1d==1.1.1   
+pip install mamba-ssm --no-cache-dir  
+https://github.com/state-spaces/mamba/issues/169 
+
+
 ## Credits
 Credits to the following repositories from which the code was adapted:
 - https://github.com/state-spaces/mamba

From b1872e1ed8542d9aee8e659f2336289a89588062 Mon Sep 17 00:00:00 2001
From: Angel Gonzalez <angegonzalez@unal.edu.co>
Date: Tue, 7 May 2024 11:35:59 +0200
Subject: [PATCH 02/33] Adding checkpoint after traning ends

---
 src/nanotron/config/config.py | 1 +
 src/nanotron/trainer.py       | 5 ++++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py
index d9946f26..e26fac75 100644
--- a/src/nanotron/config/config.py
+++ b/src/nanotron/config/config.py
@@ -129,6 +129,7 @@ class CheckpointsArgs:
     checkpoints_path: Path
     checkpoint_interval: int
     save_initial_state: Optional[bool] = False
+    save_final_state: Optional[bool] = False
     resume_checkpoint_path: Optional[Path] = None
     checkpoints_path_is_shared_file_system: Optional[bool] = False
 
diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
index 0eda00dc..70d023fb 100644
--- a/src/nanotron/trainer.py
+++ b/src/nanotron/trainer.py
@@ -442,7 +442,10 @@ def train(
                     self.save_checkpoint()
 
         dist.barrier()  # let's wait for everyone before leaving
-
+        
+        if self.config.checkpoints.save_final_state:
+            self.save_checkpoint()
+        
         self.post_training()
 
     def training_step(

From 7dd5beb36cbfbcefa2e8a67ed6186908f0e77749 Mon Sep 17 00:00:00 2001
From: Tiancheng Chen <tiachen@student.ethz.ch>
Date: Thu, 16 May 2024 18:54:55 +0200
Subject: [PATCH 03/33] fix row parallel

---
 .../parallel/tensor_parallel/functional.py    | 74 ++++++++++++++-----
 tests/test_tensor_parallel.py                 | 29 ++++++--
 2 files changed, 76 insertions(+), 27 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index fdef48ac..1c4db5de 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -387,8 +387,7 @@ def backward(ctx, grad_output):
         group = ctx.group
         use_bias = ctx.use_bias
 
-        handle_0: Optional[dist.Work] = None
-        handle_1: Optional[dist.Work] = None
+        handle: Optional[dist.Work] = None
 
         # TODO @thomasw21: gather along another dimension
         sharded_batch_size, *rest_size = grad_output.shape
@@ -412,31 +411,69 @@ def backward(ctx, grad_output):
             # https://github.com/pytorch/pytorch/blob/c47cf9bc7f9e02f649ab4ed53fe4d35732c92ab6/torch/_refs/__init__.py#L2761
             grad_output = grad_output.contiguous()
 
-            handle_0 = dist.all_gather_into_tensor(total_grad_output, grad_output, group=group, async_op=True)
-
-        grad_tensor = grad_output.matmul(weight)
+            handle = dist.all_gather_into_tensor(total_grad_output, grad_output, group=group, async_op=True)
 
-        # wait for the first all_gather to finish before starting the second all_gather
-        if handle_0 is not None:
-            handle_0.wait()
-
-        # TODO @thomasw21: gather along another dimension
-        sharded_batch_size, *rest_size = grad_tensor.shape
+        # total_grad_output: [b, s, h_out]
+        # weight: [h_out, h_in/n]
+        # total_grad_tensor: [b, s, h_in/n]
+        # grad_output: [b/n, s, h_out]
+        sharded_batch_size, *rest_size_grad_output = grad_output.shape
+        rest_size_grad_tensor = rest_size_grad_output[:-1] + [weight.shape[1]]
 
         if group.size() == 1:
-            total_grad_tensor = grad_tensor
+            total_grad_tensor = grad_output.matmul(weight)
         else:
             unsharded_batch_size = sharded_batch_size * group.size()
-
             total_grad_tensor = torch.empty(
                 unsharded_batch_size,
-                *rest_size,
-                device=grad_tensor.device,
-                dtype=grad_tensor.dtype,
+                *rest_size_grad_tensor,
+                device=grad_output.device,
+                dtype=grad_output.dtype,
                 requires_grad=False,
             )
+            before_shard_grad_tensor, same_device_shard_grad_tensor, after_shard_grad_tensor = torch.split(
+                total_grad_tensor,
+                split_size_or_sections=[
+                    sharded_batch_size * dist.get_rank(group),
+                    sharded_batch_size,
+                    sharded_batch_size * (group.size() - dist.get_rank(group) - 1),
+                ],
+                dim=0,
+            )
+            # compute local shard
+            torch.mm(
+                input=grad_output.view(-1, grad_output.shape[-1]),
+                mat2=weight,
+                out=same_device_shard_grad_tensor.view(-1, weight.shape[1]),
+            )
 
-            handle_1 = dist.all_gather_into_tensor(total_grad_tensor, grad_tensor, group=group, async_op=True)
+            if handle is not None:
+                handle.wait()
+
+            before_shard_grad_output, _, after_shard_grad_output = torch.split(
+                total_grad_output,
+                split_size_or_sections=[
+                    sharded_batch_size * dist.get_rank(group),
+                    sharded_batch_size,
+                    sharded_batch_size * (group.size() - dist.get_rank(group) - 1),
+                ],
+                dim=0,
+            )
+
+            # before shard compute
+            if before_shard_grad_tensor.numel() > 0:
+                torch.mm(
+                    input=before_shard_grad_output.view(-1, before_shard_grad_output.shape[-1]),
+                    mat2=weight,
+                    out=before_shard_grad_tensor.view(-1, weight.shape[1]),
+                )
+            # after shard compute
+            if after_shard_grad_tensor.numel() > 0:
+                torch.mm(
+                    input=after_shard_grad_output.view(-1, after_shard_grad_output.shape[-1]),
+                    mat2=weight,
+                    out=after_shard_grad_tensor.view(-1, weight.shape[1]),
+                )
 
         # Convert the tensor shapes to 2D for execution compatibility
         tensor = tensor.contiguous()
@@ -454,9 +491,6 @@ def backward(ctx, grad_output):
         grad_weight = total_grad_output.t().matmul(tensor)
         grad_bias = total_grad_output.sum(dim=0) if use_bias else None
 
-        if handle_1 is not None:
-            handle_1.wait()
-
         return total_grad_tensor, grad_weight, grad_bias, None, None
 
 
diff --git a/tests/test_tensor_parallel.py b/tests/test_tensor_parallel.py
index 127ba2fa..f5dcaeb0 100644
--- a/tests/test_tensor_parallel.py
+++ b/tests/test_tensor_parallel.py
@@ -208,14 +208,19 @@ def _test_row_linear(parallel_context: ParallelContext, tp_mode: TensorParallelL
     random_input = torch.randn(batch_size, in_features, device="cuda")
     # synchronize random_input across tp
     dist.all_reduce(random_input, op=dist.ReduceOp.AVG, group=parallel_context.tp_pg)
-
+    random_input.requires_grad = True
     # Row linear receives as input sharded input
-    random_sharded_input = random_input[
-        :,
-        dist.get_rank(parallel_context.tp_pg)
-        * in_features_per_rank : (dist.get_rank(parallel_context.tp_pg) + 1)
-        * in_features_per_rank,
-    ]
+    random_sharded_input = (
+        random_input[
+            :,
+            dist.get_rank(parallel_context.tp_pg)
+            * in_features_per_rank : (dist.get_rank(parallel_context.tp_pg) + 1)
+            * in_features_per_rank,
+        ]
+        .detach()
+        .clone()
+    )
+    random_sharded_input.requires_grad = True
 
     # Test that we get the same output after forward pass
     # TODO @kunhao: We may want to have our custom error type
@@ -261,6 +266,16 @@ def _test_row_linear(parallel_context: ParallelContext, tp_mode: TensorParallelL
     else:
         assert row_linear.bias is None
 
+    torch.testing.assert_close(
+        random_sharded_input.grad,
+        random_input.grad[
+            :,
+            dist.get_rank(parallel_context.tp_pg)
+            * in_features_per_rank : (dist.get_rank(parallel_context.tp_pg) + 1)
+            * in_features_per_rank,
+        ],
+    )
+
     parallel_context.destroy()
 
 

From ad028e6345af0342d2950da1f69ad677d8c0743d Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Fri, 31 May 2024 01:46:17 +0000
Subject: [PATCH 04/33] datatrove is all you need

---
 examples/config_nanoset.yaml                  | 24 +++---
 pyproject.toml                                |  2 +-
 run_train.py                                  |  6 +-
 src/nanotron/config/config.py                 | 16 ++--
 src/nanotron/data/collator.py                 | 80 +++++++++++++++++++
 src/nanotron/data/dataloader_builder.py       |  4 +-
 src/nanotron/data/nanoset.py                  | 75 ++++++++---------
 tests/helpers/data.py                         | 46 ++++++-----
 .../nanoset/test_build_nanoset_dataloader.py  | 79 +++++++++---------
 9 files changed, 207 insertions(+), 125 deletions(-)
 create mode 100644 src/nanotron/data/collator.py

diff --git a/examples/config_nanoset.yaml b/examples/config_nanoset.yaml
index 31f23bf0..127ddb5e 100644
--- a/examples/config_nanoset.yaml
+++ b/examples/config_nanoset.yaml
@@ -7,25 +7,25 @@ checkpoints:
 data_stages:
 - data:
     dataset:
-      dataset_path: datasets/testing_alpaca_small_input_ids.npy
+      dataset_folder: datasets/c4-es/tokenized
     num_loading_workers: 1
     seed: 42
   name: General purpose training (Single dataset)
   start_training_step: 1
 - data:
     dataset:
-      dataset_path:
-      - datasets/yelp_review_full_input_ids.npy
-      - datasets/testing_alpaca_small_input_ids.npy
+      dataset_folder:
+      - datasets/SlimPajama-6B/tokenized
+      - datasets/c4-es/tokenized
     num_loading_workers: 1
     seed: 42
   name: Second purpose training (> 1 dataset)
   start_training_step: 15
 - data:
     dataset:
-      dataset_path:
-        datasets/testing_alpaca_small_input_ids.npy: 0.8
-        datasets/yelp_review_full_input_ids.npy: 0.2
+      dataset_folder:
+        datasets/SlimPajama-6B/tokenized: 0.8
+        datasets/c4-es/tokenized: 0.2
     num_loading_workers: 1
     seed: 42
   name: Third purpose training (Blended dataset)
@@ -57,7 +57,7 @@ model:
     initializer_range: 0.02
     intermediate_size: 64
     is_llama_config: true
-    max_position_embeddings: 256
+    max_position_embeddings: 1024
     num_attention_heads: 4
     num_hidden_layers: 2
     num_key_value_heads: 4
@@ -67,7 +67,7 @@ model:
     rope_scaling: null
     tie_word_embeddings: true
     use_cache: true
-    vocab_size: 32000
+    vocab_size: 50257
 optimizer:
   accumulate_grad_in_fp32: true
   clip_grad: 1.0
@@ -88,11 +88,11 @@ optimizer:
   weight_decay: 0.01
   zero_stage: 0
 parallelism:
-  dp: 2
+  dp: 1
   expert_parallel_size: 1
   pp: 1
   pp_engine: 1f1b
-  tp: 2
+  tp: 1
   tp_linear_async_communication: true
   tp_mode: REDUCE_SCATTER
 profiler: null
@@ -105,6 +105,6 @@ tokens:
   limit_test_batches: 0
   limit_val_batches: 0
   micro_batch_size: 2
-  sequence_length: 128
+  sequence_length: 1024
   train_steps: 200
   val_check_interval: -1
diff --git a/pyproject.toml b/pyproject.toml
index e65f37a5..898d22bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,7 +49,7 @@ fast-modeling = [
 
 nanosets = [
      "transformers",
-     "datasets",
+     "datatrove[io,processing]",
      "numba",
 ]
 
diff --git a/run_train.py b/run_train.py
index b33231f4..021d955d 100644
--- a/run_train.py
+++ b/run_train.py
@@ -143,17 +143,17 @@ def get_dataloader_from_data_stage(
     elif isinstance(data.dataset, NanosetDatasetsArgs):
         # Get tokenizer cardinality
         tokenizer = AutoTokenizer.from_pretrained(trainer.config.tokenizer.tokenizer_name_or_path)
-        token_dtype = np.int32 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else np.uint16
+        token_size = 4 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else 2
         del tokenizer
         # Create Nanoset
         from nanotron.data.nanoset import Nanoset
 
         with main_rank_first(trainer.parallel_context.world_pg):
             train_dataset = Nanoset(
-                dataset_paths=data.dataset.dataset_path,
+                dataset_folders=data.dataset.dataset_folder,
                 dataset_weights=data.dataset.dataset_weights,
                 sequence_length=trainer.sequence_length,
-                token_dtype=token_dtype,
+                token_size=token_size,
                 train_split_num_samples=trainer.config.tokens.train_steps * trainer.global_batch_size,
                 random_seed=data.seed,
             )
diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py
index d5b9976f..fe194883 100644
--- a/src/nanotron/config/config.py
+++ b/src/nanotron/config/config.py
@@ -93,18 +93,18 @@ def __post_init__(self):
 
 @dataclass
 class NanosetDatasetsArgs:
-    dataset_path: Union[str, dict, List[str]]
+    dataset_folder: Union[str, dict, List[str]]
 
     def __post_init__(self):
-        if isinstance(self.dataset_path, str):  # Case 1: 1 Dataset file
-            self.dataset_path = [self.dataset_path]
+        if isinstance(self.dataset_folder, str):  # Case 1: 1 Dataset file
+            self.dataset_folder = [self.dataset_folder]
             self.dataset_weights = [1]
-        elif isinstance(self.dataset_path, List):  # Case 2: > 1 Dataset file
+        elif isinstance(self.dataset_folder, List):  # Case 2: > 1 Dataset file
             self.dataset_weights = None  # Set to None so we consume all the samples randomly
-        elif isinstance(self.dataset_path, dict):  # Case 3: dict with > 1 dataset_path and weights
-            tmp_dataset_path = self.dataset_path.copy()
-            self.dataset_path = list(tmp_dataset_path.keys())
-            self.dataset_weights = list(tmp_dataset_path.values())
+        elif isinstance(self.dataset_folder, dict):  # Case 3: dict with > 1 dataset_folder and weights
+            tmp_dataset_folder = self.dataset_folder.copy()
+            self.dataset_folder = list(tmp_dataset_folder.keys())
+            self.dataset_weights = list(tmp_dataset_folder.values())
 
 
 @dataclass
diff --git a/src/nanotron/data/collator.py b/src/nanotron/data/collator.py
new file mode 100644
index 00000000..199527e1
--- /dev/null
+++ b/src/nanotron/data/collator.py
@@ -0,0 +1,80 @@
+import dataclasses
+from typing import Dict, List, Union
+
+import numpy as np
+import torch
+from nanotron import distributed as dist
+from nanotron.parallel.context import ParallelContext
+from nanotron.parallel.pipeline_parallel.tensor_pointer import TensorPointer
+
+
+@dataclasses.dataclass
+class NanosetDataCollatorForCLM:
+    """
+    Data collator used for causal language modeling with Nanosets dataset.
+
+    - input_pp_rank: Discards last input id token
+    - output_pp_rank: Discards first label id token
+    - other pp ranks: Don't have data. Instead, we use `TensorPointer` to point to the rank having the data.
+    """
+
+    sequence_length: int
+    input_pp_rank: int
+    output_pp_rank: int
+    parallel_context: ParallelContext
+
+    def __call__(self, examples: List[Dict[str, List[np.ndarray]]]) -> Dict[str, Union[torch.Tensor, TensorPointer]]:
+        # Process the case when current rank doesn't require data. We return `TensorPointer` that points to ranks having the data.
+        current_pp_rank = dist.get_rank(self.parallel_context.pp_pg)
+        if current_pp_rank not in [
+            self.input_pp_rank,
+            self.output_pp_rank,
+        ]:
+            assert all(len(example) == 0 for example in examples)
+            return {
+                "input_ids": TensorPointer(group_rank=self.input_pp_rank),
+                "input_mask": TensorPointer(group_rank=self.input_pp_rank),
+                "label_ids": TensorPointer(group_rank=self.output_pp_rank),
+                "label_mask": TensorPointer(group_rank=self.output_pp_rank),
+            }
+
+        # Make sure we load only what's necessary, ie we only load a `input_ids` column.
+        assert all(list(example.keys()) == ["input_ids"] for example in examples)
+
+        # TODO @nouamanetazi: Is it better to have examples as np.array or torch.Tensor?
+        input_ids = torch.vstack([examples[i]["input_ids"] for i in range(len(examples))])  # (b, s)
+        batch_size, expanded_input_length = input_ids.shape
+
+        result: Dict[str, Union[torch.LongTensor, TensorPointer]] = {}
+
+        result["input_ids"] = TensorPointer(group_rank=self.input_pp_rank)
+        result["input_mask"] = TensorPointer(group_rank=self.input_pp_rank)
+        result["label_ids"] = TensorPointer(group_rank=self.output_pp_rank)
+        result["label_mask"] = TensorPointer(group_rank=self.output_pp_rank)
+
+        assert (
+            expanded_input_length == self.sequence_length + 1
+        ), f"Samples should be of length {self.sequence_length + 1} (seq_len+1), but got {expanded_input_length}"
+
+        # Process inputs: last token is the label
+        if current_pp_rank == self.input_pp_rank:
+            result["input_ids"] = input_ids[:, :-1]
+            result["input_mask"] = torch.ones((batch_size, self.sequence_length), dtype=torch.bool)
+
+        # Process labels: shift them to the left
+        if current_pp_rank == self.output_pp_rank:
+            result["label_ids"] = input_ids[:, 1:]
+            result["label_mask"] = torch.ones((batch_size, self.sequence_length), dtype=torch.bool)
+
+        if isinstance(result["input_ids"], torch.Tensor) and result["input_ids"].shape[-1] != self.sequence_length:
+            raise ValueError(
+                f"`labels` are incorrectly preprocessed. `labels` length is {result['input_ids'].shape[-1]}, but should be"
+                f" {self.sequence_length}."
+            )
+        if isinstance(result["label_ids"], torch.Tensor) and result["label_ids"].shape[-1] != self.sequence_length:
+            raise ValueError(
+                f"`labels` are incorrectly preprocessed. `labels` length is {result['label_ids'].shape[-1]}, but should be"
+                f" {self.sequence_length}."
+            )
+
+        return result
diff --git a/src/nanotron/data/dataloader_builder.py b/src/nanotron/data/dataloader_builder.py
index 4719c476..9d3285f6 100644
--- a/src/nanotron/data/dataloader_builder.py
+++ b/src/nanotron/data/dataloader_builder.py
@@ -1,7 +1,7 @@
 import nanotron.distributed as dist
 from nanotron import logging
+from nanotron.data.collator import NanosetDataCollatorForCLM
 from nanotron.dataloader import (
-    DataCollatorForCLM,
     EmptyInfiniteDataset,
     get_dataloader_worker_init,
     get_sampler,
@@ -32,7 +32,7 @@ def build_nanoset_dataloader(
         # No need to spawn a lot of workers, we can just use main
         dataloader_num_workers = 0
 
-    data_collator = DataCollatorForCLM(
+    data_collator = NanosetDataCollatorForCLM(
         sequence_length=sequence_length,
         input_pp_rank=input_pp_rank,
         output_pp_rank=output_pp_rank,
diff --git a/src/nanotron/data/nanoset.py b/src/nanotron/data/nanoset.py
index 9d62b33d..876a17e3 100644
--- a/src/nanotron/data/nanoset.py
+++ b/src/nanotron/data/nanoset.py
@@ -1,7 +1,10 @@
+import os
+import warnings
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import torch
+from datatrove.utils.dataset import DatatroveFolderDataset
 from nanotron import logging
 from nanotron.data.utils import count_dataset_indexes, normalize
 from nanotron.logging import log_rank
@@ -15,42 +18,51 @@ class Nanoset(torch.utils.data.Dataset):
     The Nanoset dataset
 
     Args:
-        dataset_paths (List[str]): List of paths to tokenized datasets
+        dataset_folders (List[str]): List of folders with tokenized datasets
         dataset_weights (List[float]): List with the weights for weighted datasets. If None, consume all samples from all datasets without weighting. Weights are normalized in __init__
         sequence_length (int): Sequence length of the built samples
-        token_dtype (Union[np.uint16, np.int32]): dtype of the tokens stored in the processed dataset files. np.uin16 for vocab sizes < 65535, np.int32 otherwise
+        token_size (int): Number of bytes for the tokens stored in the processed dataset files. 2 for vocab sizes < 65535, 4 otherwise
         train_split_num_samples (int): Number of samples the dataset needs. It's the training steps * global batch size
     """
 
     def __init__(
         self,
-        dataset_paths: List[str],
+        dataset_folders: List[str],
         dataset_weights: Union[List[float], None],
         sequence_length: int,
-        token_dtype: Union[np.uint16, np.int32],
+        token_size: int,
         train_split_num_samples: int,
         random_seed: int = 1234,
     ) -> None:
 
+        # Assertions
+        if isinstance(dataset_folders, str):
+            warnings.warn("dataset_folders should be of type List[str] but str was provided. Converting to List[str]")
+            dataset_folders = [dataset_folders]
+
         # Init
-        self.dataset_paths = dataset_paths
+        self.dataset_folders = dataset_folders
         self.dataset_weights = dataset_weights
         self.sequence_length = sequence_length
-        self.token_dtype = token_dtype
+        self.token_size = token_size
         self.train_split_num_samples = train_split_num_samples
         self.random_seed = random_seed
+        self.datatrove_datasets = []
+        for dataset_folder in self.dataset_folders:
+            self.datatrove_datasets.append(
+                DatatroveFolderDataset(
+                    folder_path=dataset_folder,
+                    filename_pattern=os.path.join(dataset_folder, "*.ds"),
+                    seq_len=sequence_length,
+                    recursive=False,
+                    token_size=token_size,
+                    shuffle=True,
+                )
+            )
 
         # Build Nanoset Index
         ## To build the index we need the length of each dataset
-        self.dataset_lengths = []
-        for dataset_path in self.dataset_paths:
-            self.dataset_buffer_mmap = np.memmap(dataset_path, mode="r", order="C", dtype=self.token_dtype)
-            self.dataset_buffer = memoryview(self.dataset_buffer_mmap)
-            dataset_number_of_tokens = int(len(self.dataset_buffer))
-            number_of_samples = int(
-                (dataset_number_of_tokens - 1) / sequence_length
-            )  # Discard last sample if length < sequence_length
-            self.dataset_lengths.append(number_of_samples)
+        self.dataset_lengths = [len(datatrove_dataset) for datatrove_dataset in self.datatrove_datasets]
         ## Set dataset weights
         if (
             self.dataset_weights is None
@@ -58,6 +70,9 @@ def __init__(
             self.dataset_weights = normalize(self.dataset_lengths)
         else:
             self.dataset_weights = normalize(dataset_weights)
+        assert len(dataset_folders) == len(
+            self.dataset_weights
+        ), f"Specified {len(self.dataset_weights)} weights but {len(dataset_folders)} datasets were provided."
         ## Build dataset index and dataset sample index
         self.dataset_index, self.dataset_sample_index = self.build_nanoset_index()
 
@@ -79,25 +94,12 @@ def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
             idx (int): The index into the dataset
 
         Returns:
-            Dict[str, numpy.ndarray]: The input ids wrapped in a dictionary
+            Dict[str, torch.LongTensor]: The input ids wrapped in a dictionary
         """
-
         dataset = self.dataset_index[idx]
         dataset_sample = self.dataset_sample_index[idx]
 
-        # Rebuild the memmap in every access to free memory
-        # https://stackoverflow.com/a/61472122
-        self.dataset_buffer_mmap = np.memmap(self.dataset_paths[dataset], mode="r", order="C", dtype=self.token_dtype)
-        self.dataset_buffer = memoryview(self.dataset_buffer_mmap)
-
-        # uint16 -> 2 bytes per token, int32 -> 4 bytes per token
-        offset = dataset_sample * self.sequence_length * (np.iinfo(self.token_dtype).bits / 8)
-        input_ids_tokens = np.frombuffer(
-            self.dataset_buffer, dtype=self.token_dtype, count=(self.sequence_length + 1), offset=int(offset)
-        )
-
-        # Return tokens as np.int32 as Torch can't handle uint16
-        return {"input_ids": input_ids_tokens.astype(np.int32)}
+        return self.datatrove_datasets[dataset][dataset_sample]
 
     def build_nanoset_index(self) -> np.ndarray:
         """
@@ -124,15 +126,6 @@ def build_nanoset_index(self) -> np.ndarray:
 
         return dataset_index, dataset_sample_index
 
-    def __del__(self) -> None:
-        """
-        Clean up Nanoset
-        """
-
-        if hasattr(self, "dataset_buffer_mmap"):
-            self.dataset_buffer_mmap._mmap.close()
-        del self.dataset_buffer_mmap
-
     def print_nanoset_info(self):
 
         log_rank(f"> Total number of samples: {len(self)}", logger=logger, level=logging.INFO, rank=0)
@@ -141,10 +134,10 @@ def print_nanoset_info(self):
         )
 
         # Print samples from each dataset + weight
-        dataset_sample_count = count_dataset_indexes(self.dataset_index, len(self.dataset_paths))
+        dataset_sample_count = count_dataset_indexes(self.dataset_index, len(self.dataset_folders))
         for index, sample_count in enumerate(dataset_sample_count):
             log_rank(
-                f">   Total number of samples from the {self.dataset_paths[index].rsplit('/', 1)[-1]} dataset: {sample_count} ({round(normalize(dataset_sample_count).tolist()[index], 2)})",
+                f">   Total number of samples from the {self.dataset_folders[index]} dataset: {sample_count} ({round(normalize(dataset_sample_count).tolist()[index], 2)})",
                 logger=logger,
                 level=logging.INFO,
                 rank=0,
diff --git a/tests/helpers/data.py b/tests/helpers/data.py
index 33bb2480..d01c717f 100644
--- a/tests/helpers/data.py
+++ b/tests/helpers/data.py
@@ -10,46 +10,53 @@
 package_path = Path(package.__file__).parent.parent.parent
 sys.path.append(str(package_path))
 
-from argparse import Namespace
 
 import nanotron.distributed as dist
 import torch
+from datatrove.executor.local import LocalPipelineExecutor
+from datatrove.pipeline.readers import JsonlReader
+from datatrove.pipeline.tokens.tokenizer import DocumentTokenizer
 from nanotron.data.nanoset import Nanoset
 from nanotron.parallel import ParallelContext
 from nanotron.parallel.pipeline_parallel.tensor_pointer import TensorPointer
 from nanotron.sanity_checks import assert_tensor_synced_across_pg
 
-from tools.preprocess_data import main
-
 
 def create_dataset_paths(tmp_dir: str, quantity: int):
-    json_dataset_path = [os.path.join(tmp_dir, f"pytest_{i}") for i in range(quantity)]
-    mmap_dataset_path = [f"{path}_input_ids.npy" for path in json_dataset_path]
+    json_dataset_path = [os.path.join(tmp_dir, f"pytest_{i}.json") for i in range(quantity)]
+    datatrove_tokenized_dataset_paths = [os.path.join(tmp_dir, f"tokenized_documents_{i}") for i in range(quantity)]
 
-    return json_dataset_path, mmap_dataset_path
+    return json_dataset_path, datatrove_tokenized_dataset_paths
 
 
 def create_dummy_json_dataset(path_to_json: str, dummy_text: str, n_samples: int = 50000):
 
-    with open(path_to_json + ".json", "a") as json_file:
+    with open(path_to_json, "a") as json_file:
         for sample in range(n_samples):
             sample_dict = {"text": f"[{sample}] Hello! Im sample {sample}! And this is my dummy text: {dummy_text}"}
             json_file.write(json.dumps(sample_dict))
             json_file.write("\n")
 
 
-def preprocess_dummy_dataset(path_to_json: str, tokenizer: str):
-    # Create args for preprocessing
-    args = Namespace(
-        input=path_to_json + ".json",
-        column="text",
-        output_prefix=path_to_json,
-        tokenizer_name_or_path=tokenizer,
-        add_special_tokens=False,
+def preprocess_dummy_dataset(json_dataset_path: str, datatrove_tokenized_dataset_path: str, tokenizer: str):
+    tmp_dir = str(Path(json_dataset_path).parent.absolute())
+
+    # Datatrove tokenizing pipeline
+    dist_executor = LocalPipelineExecutor(
+        pipeline=[
+            JsonlReader(data_folder=json_dataset_path),
+            DocumentTokenizer(
+                output_folder=datatrove_tokenized_dataset_path,
+                local_working_dir=tmp_dir,
+                save_filename="dummy_dataset_tokenized",
+                tokenizer_name_or_path=tokenizer,
+                eos_token=None,
+            ),
+        ],
+        tasks=1,
+        workers=-1,
     )
-
-    # tools/preprocess_data.py main
-    main(args)
+    dist_executor.run()
 
 
 def assert_batch_dataloader(
@@ -122,7 +129,7 @@ def assert_nanoset_sync_across_all_ranks(nanoset: Nanoset, parallel_context: Par
     IDX_SAMPLE = 23
 
     nanoset_identifiers = OrderedDict()
-    nanoset_identifiers["dataset_paths"] = nanoset.dataset_paths
+    nanoset_identifiers["dataset_folders"] = nanoset.dataset_folders
     nanoset_identifiers["dataset_weights"] = nanoset.dataset_weights.tolist()
     nanoset_identifiers["sequence_length"] = nanoset.sequence_length
     nanoset_identifiers["train_split_num_samples"] = nanoset.train_split_num_samples
@@ -131,6 +138,7 @@ def assert_nanoset_sync_across_all_ranks(nanoset: Nanoset, parallel_context: Par
     nanoset_identifiers["input_ids"] = nanoset[IDX_SAMPLE]["input_ids"].tolist()
     nanoset_identifiers["dataset_index"] = nanoset.dataset_index.tolist()
     nanoset_identifiers["dataset_sample_index"] = nanoset.dataset_sample_index.tolist()
+    nanoset_identifiers["token_size"] = nanoset.token_size
 
     unique_description_hash = compute_hash(nanoset_identifiers)
     assert_tensor_synced_across_pg(
diff --git a/tests/nanoset/test_build_nanoset_dataloader.py b/tests/nanoset/test_build_nanoset_dataloader.py
index 2c3ff542..113c545c 100644
--- a/tests/nanoset/test_build_nanoset_dataloader.py
+++ b/tests/nanoset/test_build_nanoset_dataloader.py
@@ -1,6 +1,7 @@
 import sys
 from math import isclose
 from pathlib import Path
+from typing import List
 
 package_path = Path(__file__).parent.parent
 sys.path.append(str(package_path))
@@ -33,7 +34,7 @@
         for all_3d_configs in get_all_3d_configurations(gpus)
     ],
 )
-@pytest.mark.parametrize("train_steps", [5, 100])
+@pytest.mark.parametrize("train_steps", [500, 10000])
 @pytest.mark.parametrize("sequence_length", [512, 8192])
 @pytest.mark.parametrize("tokenizer_name_or_path", ["openai-community/gpt2", "unsloth/llama-3-8b-bnb-4bit"])
 @rerun_if_address_is_in_use()
@@ -42,16 +43,21 @@ def test_build_nanoset_dataloader(
 ):
     test_context = TestContext()
 
-    # Create dataset files
-    json_paths, mmap_dataset_paths = create_dataset_paths(tmp_dir=test_context.get_auto_remove_tmp_dir(), quantity=2)
+    # Create dataset folders
+    json_paths, datatrove_tokenized_dataset_folders = create_dataset_paths(
+        tmp_dir=test_context.get_auto_remove_tmp_dir(), quantity=2
+    )
 
     # Create dummy json datasets
     for idx, json_path in enumerate(json_paths):
         create_dummy_json_dataset(path_to_json=json_path, dummy_text=f"Nanoset {idx}!", n_samples=(idx + 1) * 50000)
 
+    # Preprocess json dataset with datatrove
+    for json_path, datatrove_tokenized_dataset_folder in zip(json_paths, datatrove_tokenized_dataset_folders):
+        preprocess_dummy_dataset(json_path, datatrove_tokenized_dataset_folder, tokenizer_name_or_path)
+
     init_distributed(tp=tp, dp=dp, pp=pp)(_test_build_nanoset_dataloader)(
-        json_paths=json_paths,
-        path_to_mmap_files=mmap_dataset_paths,
+        datatrove_tokenized_dataset_folders=datatrove_tokenized_dataset_folders,
         train_steps=train_steps,
         sequence_length=sequence_length,
         tokenizer_name_or_path=tokenizer_name_or_path,
@@ -60,8 +66,7 @@ def test_build_nanoset_dataloader(
 
 def _test_build_nanoset_dataloader(
     parallel_context: ParallelContext,
-    json_paths: str,
-    path_to_mmap_files: str,
+    datatrove_tokenized_dataset_folders: List[str],
     train_steps: int,
     sequence_length: int,
     tokenizer_name_or_path: str,
@@ -71,41 +76,37 @@ def _test_build_nanoset_dataloader(
     N_MICRO_BATCHES_PER_BATCH = 8
     GLOBAL_BATCH_SIZE = MICRO_BATCH_SIZE * N_MICRO_BATCHES_PER_BATCH * parallel_context.dp_pg.size()
 
-    # Preprocess dummy json datasets
-    for json_path in json_paths:
-        preprocess_dummy_dataset(path_to_json=json_path, tokenizer=tokenizer_name_or_path)
-
     input_pp_rank, output_pp_rank = 0, int(parallel_context.pp_pg.size() - 1)
 
     # Get tokenizer cardinality
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
-    token_dtype = np.int32 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else np.uint16
+    token_size = 4 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else 2
     del tokenizer
 
     # Create Nanoset configs: 1. Normal 2. Blended 3. Blended with weights
     nanoset_config = {
-        "dataset_paths": [path_to_mmap_files[0]],
+        "dataset_folders": [datatrove_tokenized_dataset_folders[0]],
         "dataset_weights": [1],
         "sequence_length": sequence_length,
-        "token_dtype": token_dtype,
+        "token_size": token_size,
         "train_split_num_samples": train_steps * GLOBAL_BATCH_SIZE,
         "random_seed": SEED,
     }
 
     blended_nanoset_config = {
-        "dataset_paths": [path_to_mmap_files[0], path_to_mmap_files[1]],
+        "dataset_folders": datatrove_tokenized_dataset_folders,
         "dataset_weights": None,
         "sequence_length": sequence_length,
-        "token_dtype": token_dtype,
+        "token_size": token_size,
         "train_split_num_samples": train_steps * GLOBAL_BATCH_SIZE,
         "random_seed": SEED,
     }
 
     blended_weighted_nanoset_config = {
-        "dataset_paths": [path_to_mmap_files[0], path_to_mmap_files[1]],
+        "dataset_folders": datatrove_tokenized_dataset_folders,
         "dataset_weights": [8, 2],
         "sequence_length": sequence_length,
-        "token_dtype": token_dtype,
+        "token_size": token_size,
         "train_split_num_samples": train_steps * GLOBAL_BATCH_SIZE,
         "random_seed": SEED,
     }
@@ -119,7 +120,7 @@ def _test_build_nanoset_dataloader(
 
         # Assert we have the same Nanoset in all ranks
         assert_nanoset_sync_across_all_ranks(train_dataset, parallel_context)
-        dataset_sample_count = count_dataset_indexes(train_dataset.dataset_index, len(train_dataset.dataset_paths))
+        dataset_sample_count = count_dataset_indexes(train_dataset.dataset_index, len(train_dataset.dataset_folders))
         for idx, ds_length in enumerate(train_dataset.dataset_lengths):
             # Assert Nanoset doesn't sample indexes greater than the datasets
             assert (
@@ -129,7 +130,7 @@ def _test_build_nanoset_dataloader(
             # Assert Nanoset builds up the correct blend WRT the dataset_weights
             assert isclose(
                 normalize(dataset_sample_count).tolist()[idx], train_dataset.dataset_weights[idx], abs_tol=0.05
-            ), f"Requested Nanoset to contain {round(train_dataset.dataset_weights[idx]*100, 2)}% of samples from {train_dataset.dataset_paths[idx]} but got {round(normalize(dataset_sample_count).tolist()[idx]*100, 2)}%"
+            ), f"Requested Nanoset to contain {round(train_dataset.dataset_weights[idx]*100, 2)}% of samples from {train_dataset.dataset_folders[idx]} but got {round(normalize(dataset_sample_count).tolist()[idx]*100, 2)}%"
         # Create Dataloaders
         dataloader = build_nanoset_dataloader(
             train_dataset,
@@ -162,22 +163,27 @@ def _test_build_nanoset_dataloader(
         for all_3d_configs in get_all_3d_configurations(gpus)
     ],
 )
-@pytest.mark.parametrize("skipped_batches", [20, 50])
+@pytest.mark.parametrize("skipped_batches", [20, 5555])
 @pytest.mark.parametrize("tokenizer_name_or_path", ["openai-community/gpt2", "unsloth/llama-3-8b-bnb-4bit"])
 @rerun_if_address_is_in_use()
 def test_recover_nanoset_dataloader(tp: int, dp: int, pp: int, skipped_batches: int, tokenizer_name_or_path: str):
     test_context = TestContext()
 
-    # Create dataset files
-    json_paths, mmap_dataset_paths = create_dataset_paths(tmp_dir=test_context.get_auto_remove_tmp_dir(), quantity=2)
+    # Create dataset folders
+    json_paths, datatrove_tokenized_dataset_folders = create_dataset_paths(
+        tmp_dir=test_context.get_auto_remove_tmp_dir(), quantity=2
+    )
 
     # Create dummy json datasets
     for idx, json_path in enumerate(json_paths):
         create_dummy_json_dataset(path_to_json=json_path, dummy_text=f"Nanoset {idx}!", n_samples=(idx + 1) * 50000)
 
+    # Preprocess json dataset with datatrove
+    for json_path, datatrove_tokenized_dataset_folder in zip(json_paths, datatrove_tokenized_dataset_folders):
+        preprocess_dummy_dataset(json_path, datatrove_tokenized_dataset_folder, tokenizer_name_or_path)
+
     init_distributed(tp=tp, dp=dp, pp=pp)(_test_recover_nanoset_dataloader)(
-        json_paths=json_paths,
-        path_to_mmap_files=mmap_dataset_paths,
+        datatrove_tokenized_dataset_folders=datatrove_tokenized_dataset_folders,
         skipped_batches=skipped_batches,
         tokenizer_name_or_path=tokenizer_name_or_path,
     )
@@ -185,8 +191,7 @@ def test_recover_nanoset_dataloader(tp: int, dp: int, pp: int, skipped_batches:
 
 def _test_recover_nanoset_dataloader(
     parallel_context: ParallelContext,
-    json_paths: str,
-    path_to_mmap_files: str,
+    datatrove_tokenized_dataset_folders: List[str],
     skipped_batches: int,
     tokenizer_name_or_path: str,
 ):
@@ -195,43 +200,39 @@ def _test_recover_nanoset_dataloader(
     N_MICRO_BATCHES_PER_BATCH = 8
     GLOBAL_BATCH_SIZE = MICRO_BATCH_SIZE * N_MICRO_BATCHES_PER_BATCH * parallel_context.dp_pg.size()
     SEQUENCE_LENGTH = 1024
-    TRAIN_STEPS = 100
-
-    # Preprocess dummy json datasets
-    for json_path in json_paths:
-        preprocess_dummy_dataset(path_to_json=json_path, tokenizer=tokenizer_name_or_path)
+    TRAIN_STEPS = 10000
 
     input_pp_rank, output_pp_rank = 0, int(parallel_context.pp_pg.size() - 1)
 
     # Get tokenizer cardinality
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
-    token_dtype = np.int32 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else np.uint16
+    token_size = 4 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else 2
     del tokenizer
 
     # Create Nanoset configs: 1. Normal 2. Blended 3. Blended with weights
     nanoset_config = {
-        "dataset_paths": [path_to_mmap_files[0]],
+        "dataset_folders": [datatrove_tokenized_dataset_folders[0]],
         "dataset_weights": [1],
         "sequence_length": SEQUENCE_LENGTH,
-        "token_dtype": token_dtype,
+        "token_size": token_size,
         "train_split_num_samples": TRAIN_STEPS * GLOBAL_BATCH_SIZE,
         "random_seed": SEED,
     }
 
     blended_nanoset_config = {
-        "dataset_paths": [path_to_mmap_files[0], path_to_mmap_files[1]],
+        "dataset_folders": datatrove_tokenized_dataset_folders,
         "dataset_weights": None,
         "sequence_length": SEQUENCE_LENGTH,
-        "token_dtype": token_dtype,
+        "token_size": token_size,
         "train_split_num_samples": TRAIN_STEPS * GLOBAL_BATCH_SIZE,
         "random_seed": SEED,
     }
 
     blended_weighted_nanoset_config = {
-        "dataset_paths": [path_to_mmap_files[0], path_to_mmap_files[1]],
+        "dataset_folders": datatrove_tokenized_dataset_folders,
         "dataset_weights": [8, 2],
         "sequence_length": SEQUENCE_LENGTH,
-        "token_dtype": token_dtype,
+        "token_size": token_size,
         "train_split_num_samples": TRAIN_STEPS * GLOBAL_BATCH_SIZE,
         "random_seed": SEED,
     }

From df1632b260a47b3ad5678ee9b375c7055c8c9848 Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Fri, 31 May 2024 18:17:56 +0000
Subject: [PATCH 05/33] Refractored preprocessing script to work with datatrove

---
 tools/preprocess_data.py | 132 +++++++++++++++------------------------
 1 file changed, 49 insertions(+), 83 deletions(-)

diff --git a/tools/preprocess_data.py b/tools/preprocess_data.py
index 465d22f0..75dfd94a 100644
--- a/tools/preprocess_data.py
+++ b/tools/preprocess_data.py
@@ -1,26 +1,27 @@
 import argparse
-import os
-import shutil
-import sys
 
-import numpy as np
-import torch.distributed as dist
-from tqdm import tqdm
-from transformers import AutoTokenizer
-
-from datasets import concatenate_datasets, load_dataset
+from datatrove.executor.local import LocalPipelineExecutor
+from datatrove.pipeline.readers import HuggingFaceDatasetReader
+from datatrove.pipeline.tokens import DocumentTokenizer
 
 
 def get_args():
     parser = argparse.ArgumentParser()
-    group = parser.add_argument_group(title="input data")
+    group = parser.add_argument_group(title="Dataset reader")
     group.add_argument(
-        "--input", type=str, required=True, help="Path to local stored dataset or repository on the Hugging Face hub"
+        "--dataset",
+        type=str,
+        required=True,
+        help="Path to local stored dataset or repository on the Hugging Face hub that can be loaded with datasets.load_dataset",
+    )
+    group.add_argument(
+        "--column", type=str, default="text", help="Column to preprocess from the Dataset. Default: text"
+    )
+    parser.add_argument(
+        "--split", type=str, default="train", help="Which split of the data to process. Default: train"
     )
-    group.add_argument("--column", type=str, default="text", help="Column to preprocess from the Dataset")
-    parser.add_argument("--split", type=str, default="train", help="Which split of the data to process")
 
-    group = parser.add_argument_group(title="tokenizer")
+    group = parser.add_argument_group(title="Tokenizer")
     group.add_argument(
         "--tokenizer-name-or-path",
         type=str,
@@ -28,13 +29,26 @@ def get_args():
         help="A path to a directory containing vocabulary files required by the tokenizer or the model id of a predefined tokenizer hosted inside a model repo on the Hugging Face Hub.",
     )
     group.add_argument(
-        "--add-special-tokens",
-        action="store_true",
-        help="Whether or not to add special tokens when encoding the sequences. This will be passed to the Tokenizer",
+        "--eos-token",
+        type=str,
+        default=None,
+        help="EOS token to add after each document. Default: None",
     )
 
-    group = parser.add_argument_group(title="output data")
-    group.add_argument("--output-prefix", type=str, required=True, help="Path to the output processed dataset file")
+    group = parser.add_argument_group(title="Output data")
+    group.add_argument(
+        "--output-folder", type=str, required=True, help="Path to the output folder to store the tokenized documents"
+    )
+    group = parser.add_argument_group(title="Miscellaneous configs")
+    group.add_argument(
+        "--logging-dir",
+        type=str,
+        default=None,
+        help="Path to a folder for storing the logs of the preprocessing step. Default: None",
+    )
+    group.add_argument(
+        "--n-tasks", type=int, default=8, help="Total number of tasks to run the preprocessing step. Default: 8"
+    )
 
     args = parser.parse_args()
 
@@ -43,73 +57,25 @@ def get_args():
 
 def main(args):
 
-    world_size, rank = int(os.environ["WORLD_SIZE"]), int(os.environ["RANK"])
-
-    # Remove stdout from all processes except main to not flood the stdout
-    if rank:
-        sys.stdout = open(os.devnull, "w")
-
-    # Check if output directory exists
-    if not os.path.isdir(os.path.abspath(os.path.join(args.output_prefix, os.path.pardir))):
-        print(f"Creating {os.path.abspath(os.path.join(args.output_prefix, os.path.pardir))} directory...")
-        os.makedirs(os.path.abspath(os.path.join(args.output_prefix, os.path.pardir)), exist_ok=True)
-
-    if args.input.endswith(".json"):  # For processing JSON files (Cross compatibility with other projects)
-        ds = load_dataset("json", data_files=args.input)
-        ds = concatenate_datasets(
-            [ds[splits] for splits in ds.keys()]
-        )  # load_dataset returns DatasetDict and we want a Dataset
-    else:
-        ds = load_dataset(args.input, split=args.split)
-
-    ds = ds.shard(num_shards=world_size, index=rank, contiguous=True)
-    ds = ds.select_columns(args.column)
-
-    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name_or_path)
-    token_dtype = np.int32 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else np.uint16
-
-    # Create tmp directory for worker outputs
-    tmp_folder = os.path.abspath(os.path.join(args.output_prefix, os.pardir, "tmp"))
-    os.makedirs(tmp_folder, exist_ok=True)
-
-    print("Creating workers output files...")
-    worker_output_file = os.path.join(tmp_folder, f"worker_{rank}_input_ids.npy")
-    ds = ds.map(
-        lambda x: {"input_ids": tokenizer(x, add_special_tokens=args.add_special_tokens).input_ids},
-        input_columns=args.column,
-        batched=True,
-        desc="Tokenizing Dataset",
-        remove_columns=[args.column],
+    preprocess_executor = LocalPipelineExecutor(
+        pipeline=[
+            HuggingFaceDatasetReader(
+                dataset=args.dataset,
+                text_key=args.column,
+                dataset_options={"split": args.split},
+            ),
+            DocumentTokenizer(
+                output_folder=args.output_folder,
+                tokenizer_name_or_path=args.tokenizer_name_or_path,
+                eos_token=args.eos_token,
+            ),
+        ],
+        tasks=args.n_tasks,
+        logging_dir=args.logging_dir,
     )
-
-    worker_input_ids_file = open(worker_output_file, "wb")
-    for sample in ds:
-        np_array = np.array(sample["input_ids"], dtype=token_dtype)
-        worker_input_ids_file.write(np_array.tobytes(order="C"))
-    worker_input_ids_file.close()
-
-    # Wait for all workers to process each shard of the Dataset
-    dist.barrier()
-
-    # Only the main rank merges the worker files
-    if not rank:
-        output_file = f"{args.output_prefix}_input_ids.npy"
-        input_ids_file = open(output_file, "wb")
-        for worker_idx in tqdm(range(world_size), desc="Merging workers output files"):
-            worker_output_file = os.path.join(tmp_folder, f"worker_{worker_idx}_input_ids.npy")
-            with open(worker_output_file, "rb") as f:
-                shutil.copyfileobj(f, input_ids_file)
-            os.remove(worker_output_file)
-
-        input_ids_file.close()
-        os.rmdir(tmp_folder)
-        print(f"Done! {args.input} processed dataset stored in {output_file}")
-
-    else:  # Close devnull stdout redirect
-        sys.stdout.close()
+    preprocess_executor.run()
 
 
 if __name__ == "__main__":
     _args = get_args()
-    dist.init_process_group(backend="gloo")
     main(_args)

From 395a4dbbd9124d8eb90db66772bd702e8eeff5d9 Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Sat, 1 Jun 2024 15:21:12 +0000
Subject: [PATCH 06/33] Updated dataset_weights default value

---
 src/nanotron/data/nanoset.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/nanotron/data/nanoset.py b/src/nanotron/data/nanoset.py
index 876a17e3..90200967 100644
--- a/src/nanotron/data/nanoset.py
+++ b/src/nanotron/data/nanoset.py
@@ -19,7 +19,7 @@ class Nanoset(torch.utils.data.Dataset):
 
     Args:
         dataset_folders (List[str]): List of folders with tokenized datasets
-        dataset_weights (List[float]): List with the weights for weighted datasets. If None, consume all samples from all datasets without weighting. Weights are normalized in __init__
+        dataset_weights (Union[List[float], None]): List with the weights for weighted datasets. If None, consume all samples from all datasets without weighting. Weights are normalized in __init__
         sequence_length (int): Sequence length of the built samples
         token_size (int): Number of bytes for the tokens stored in the processed dataset files. 2 for vocab sizes < 65535, 4 otherwise
         train_split_num_samples (int): Number of samples the dataset needs. It's the training steps * global batch size
@@ -28,21 +28,20 @@ class Nanoset(torch.utils.data.Dataset):
     def __init__(
         self,
         dataset_folders: List[str],
-        dataset_weights: Union[List[float], None],
         sequence_length: int,
         token_size: int,
         train_split_num_samples: int,
+        dataset_weights: Union[List[float], None] = None,
         random_seed: int = 1234,
     ) -> None:
 
-        # Assertions
+        # Checks
         if isinstance(dataset_folders, str):
             warnings.warn("dataset_folders should be of type List[str] but str was provided. Converting to List[str]")
             dataset_folders = [dataset_folders]
 
         # Init
         self.dataset_folders = dataset_folders
-        self.dataset_weights = dataset_weights
         self.sequence_length = sequence_length
         self.token_size = token_size
         self.train_split_num_samples = train_split_num_samples
@@ -65,7 +64,7 @@ def __init__(
         self.dataset_lengths = [len(datatrove_dataset) for datatrove_dataset in self.datatrove_datasets]
         ## Set dataset weights
         if (
-            self.dataset_weights is None
+            dataset_weights is None
         ):  # Case of training with > 1 datasets without weighting them: Consume both datasets entirely on each epoch
             self.dataset_weights = normalize(self.dataset_lengths)
         else:

From 5f8a52b08b702e206f31f2660e4b6f22ac328c95 Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Sat, 1 Jun 2024 15:24:36 +0000
Subject: [PATCH 07/33] Support for jsonl files for preprocess_data.py and
 updated preprocessing step in tests

---
 tests/helpers/data.py    | 37 ++++++++++------------
 tools/preprocess_data.py | 66 ++++++++++++++++++++++++++++------------
 2 files changed, 63 insertions(+), 40 deletions(-)

diff --git a/tests/helpers/data.py b/tests/helpers/data.py
index d01c717f..72deb7f5 100644
--- a/tests/helpers/data.py
+++ b/tests/helpers/data.py
@@ -3,6 +3,7 @@
 import json
 import os
 import sys
+from argparse import Namespace
 from collections import OrderedDict
 from pathlib import Path
 
@@ -10,17 +11,15 @@
 package_path = Path(package.__file__).parent.parent.parent
 sys.path.append(str(package_path))
 
-
 import nanotron.distributed as dist
 import torch
-from datatrove.executor.local import LocalPipelineExecutor
-from datatrove.pipeline.readers import JsonlReader
-from datatrove.pipeline.tokens.tokenizer import DocumentTokenizer
 from nanotron.data.nanoset import Nanoset
 from nanotron.parallel import ParallelContext
 from nanotron.parallel.pipeline_parallel.tensor_pointer import TensorPointer
 from nanotron.sanity_checks import assert_tensor_synced_across_pg
 
+from tools.preprocess_data import main
+
 
 def create_dataset_paths(tmp_dir: str, quantity: int):
     json_dataset_path = [os.path.join(tmp_dir, f"pytest_{i}.json") for i in range(quantity)]
@@ -39,24 +38,20 @@ def create_dummy_json_dataset(path_to_json: str, dummy_text: str, n_samples: int
 
 
 def preprocess_dummy_dataset(json_dataset_path: str, datatrove_tokenized_dataset_path: str, tokenizer: str):
-    tmp_dir = str(Path(json_dataset_path).parent.absolute())
-
-    # Datatrove tokenizing pipeline
-    dist_executor = LocalPipelineExecutor(
-        pipeline=[
-            JsonlReader(data_folder=json_dataset_path),
-            DocumentTokenizer(
-                output_folder=datatrove_tokenized_dataset_path,
-                local_working_dir=tmp_dir,
-                save_filename="dummy_dataset_tokenized",
-                tokenizer_name_or_path=tokenizer,
-                eos_token=None,
-            ),
-        ],
-        tasks=1,
-        workers=-1,
+    # Create args for preprocessing
+    args = Namespace(
+        readers="jsonl",
+        dataset=json_dataset_path,
+        column="text",
+        glob_pattern=None,
+        output_folder=datatrove_tokenized_dataset_path,
+        tokenizer_name_or_path=tokenizer,
+        eos_token=None,
+        n_tasks=1,
+        logging_dir=None,
     )
-    dist_executor.run()
+    # tools/preprocess_data.py main
+    main(args)
 
 
 def assert_batch_dataloader(
diff --git a/tools/preprocess_data.py b/tools/preprocess_data.py
index 75dfd94a..38db67f1 100644
--- a/tools/preprocess_data.py
+++ b/tools/preprocess_data.py
@@ -1,25 +1,19 @@
+"""
+To process HuggingFace Datasets:
+    python3 tools/preprocess_data.py --tokenizer-name-or-path meta-llama/Meta-Llama-3-8B --output-folder datasets/emotion --n-tasks 16 hf --dataset dair-ai/emotion
+To process Jsonl files:
+    python3 tools/preprocess_data.py --tokenizer-name-or-path meta-llama/Meta-Llama-3-8B --output-folder datasets/c4-es --n-tasks 16 jsonl --dataset raw_datasets/c4-es-json-files
+"""
+
 import argparse
 
 from datatrove.executor.local import LocalPipelineExecutor
-from datatrove.pipeline.readers import HuggingFaceDatasetReader
+from datatrove.pipeline.readers import HuggingFaceDatasetReader, JsonlReader
 from datatrove.pipeline.tokens import DocumentTokenizer
 
 
 def get_args():
     parser = argparse.ArgumentParser()
-    group = parser.add_argument_group(title="Dataset reader")
-    group.add_argument(
-        "--dataset",
-        type=str,
-        required=True,
-        help="Path to local stored dataset or repository on the Hugging Face hub that can be loaded with datasets.load_dataset",
-    )
-    group.add_argument(
-        "--column", type=str, default="text", help="Column to preprocess from the Dataset. Default: text"
-    )
-    parser.add_argument(
-        "--split", type=str, default="train", help="Which split of the data to process. Default: train"
-    )
 
     group = parser.add_argument_group(title="Tokenizer")
     group.add_argument(
@@ -49,6 +43,34 @@ def get_args():
     group.add_argument(
         "--n-tasks", type=int, default=8, help="Total number of tasks to run the preprocessing step. Default: 8"
     )
+    # Subparsers for processing either Hugging Face datasets or jsonl files
+    sp = parser.add_subparsers(
+        dest="readers",
+        required=True,
+        description="Type of dataset to process. It can be either a Hugging Face Dataset loaded with datasets.load_data ('hf') or a .jsonl dataset ('jsonl')",
+    )
+
+    p1 = sp.add_parser(name="hf")
+    p1.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help="Path to local stored dataset or repository on the Hugging Face hub that can be loaded with datasets.load_dataset",
+    )
+    p1.add_argument("--column", type=str, default="text", help="Column to preprocess from the Dataset. Default: text")
+    p1.add_argument("--split", type=str, default="train", help="Which split of the data to process. Default: train")
+
+    p2 = sp.add_parser(name="jsonl")
+    p2.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help="Path to a .jsonl file or a folder containing multiple .jsonl files",
+    )
+    p2.add_argument("--column", type=str, default="text", help="Column to preprocess from the Dataset. Default: text")
+    p2.add_argument(
+        "--glob-pattern", type=str, default=None, help="A glob pattern to filter files to read. Default: None"
+    )
 
     args = parser.parse_args()
 
@@ -56,18 +78,24 @@ def get_args():
 
 
 def main(args):
+    # Build datatrove reader
+    if args.readers == "hf":
+        datatrove_reader = HuggingFaceDatasetReader(
+            dataset=args.dataset,
+            text_key=args.column,
+            dataset_options={"split": args.split},
+        )
+    else:
+        datatrove_reader = JsonlReader(data_folder=args.dataset, text_key=args.column, glob_pattern=args.glob_pattern)
 
     preprocess_executor = LocalPipelineExecutor(
         pipeline=[
-            HuggingFaceDatasetReader(
-                dataset=args.dataset,
-                text_key=args.column,
-                dataset_options={"split": args.split},
-            ),
+            datatrove_reader,
             DocumentTokenizer(
                 output_folder=args.output_folder,
                 tokenizer_name_or_path=args.tokenizer_name_or_path,
                 eos_token=args.eos_token,
+                max_tokens_per_file=1e9,
             ),
         ],
         tasks=args.n_tasks,

From 9fcd071035b45b6a2cdf030142c6c0b7760bebf1 Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Sat, 1 Jun 2024 15:26:06 +0000
Subject: [PATCH 08/33] Updated nanoset docs

---
 docs/nanoset.md | 60 ++++++++++++++++++++++++++-----------------------
 1 file changed, 32 insertions(+), 28 deletions(-)

diff --git a/docs/nanoset.md b/docs/nanoset.md
index 02649bd0..9dce21b7 100644
--- a/docs/nanoset.md
+++ b/docs/nanoset.md
@@ -1,41 +1,42 @@
 # Nanosets
-Nanotron incorporates [`Nanosets`](../src/nanotron/data/nanoset.py), a kind of datasets based on [numpy memory-mapped arrays](https://numpy.org/doc/stable/reference/generated/numpy.memmap.html). `Nanosets` are capable of serving batches from files containing pre-tokenized datasets. They allow reading tokens from one or multiple datasets and even specifying the weight of each dataset when building batches.
+Nanotron incorporates [`Nanosets`](../src/nanotron/data/nanoset.py), a dataset for processing tokenized documents with [`datatrove`](https://github.com/huggingface/datatrove). They allow reading tokens from one or multiple datasets and even specifying the weight of each dataset when building batches.
 ## Install
 To use `Nanosets`, it's necessary to install Nanotron with the `nanosets` flavor.
 ```
-pip install -e '.[nanosets]'
+pip install nanotron[nanosets]
 ```
 This will install the following dependencies:
-- `transformers`: To tokenize the datasets
-- `datasets`: To preprocess the datasets
+- `datatrove`: To preprocess the datasets
 - `numba`: To compile helper functions in order to speed up the creation of `Nanosets`
+- `transformers`: For the tokenizers
 ## Data pre-processing
-To use these datasets, first, we need to preprocess the data. The input format can either be a column of a Hugging Face Dataset or a .json file containing a text sample per line. For example:
+To use this dataset, first, we need to preprocess the data using `datatrove`'s `DocumentTokenizer` pipeline. We invite you to take a look at `datatrove`, since it contains multiple features that allow, for example, filter out documents based on specific rules/criteria, extract text content from raw formats or scheduling the preprocessing in a Slurm cluster. We have also added a simple script capable of tokenizing datasets.
 
-<pre>
-{"src": "www.nvidia.com", "text": "The quick brown fox", "type": "Eng", "id": "0", "title": "First Part"}
-{"src": "The Internet", "text": "jumps over the lazy dog", "type": "Eng", "id": "42", "title": "Second Part"}
-</pre>
-
-The preprocessing is done using the [`tools/preprocess_data.py`](../tools/preprocess_data.py) script. Below we show an example for processing a corpus with the Llama2 tokenizer.
+The preprocessing is done using the [`tools/preprocess_data.py`](../tools/preprocess_data.py) script. The input format can either be a Hugging Face Dataset, a path to a `.jsonl` or a path to a folder containing multiple `.jsonl` files. Below we show an example for processing a Hugging Face Dataset from the Hub with the Llama3 tokenizer.
 
 <pre>
-torchrun --nproc-per-node 16 tools/preprocess_data.py \
-       --input HuggingFaceH4/testing_alpaca_small \
-       --split train \
-       --column completion \
-       --output-prefix datasets/testing_alpaca_small \
-       --tokenizer-name-or-path openai-community/gpt2
+python3 tools/preprocess_data.py \
+       --tokenizer-name-or-path meta-llama/Meta-Llama-3-8B \
+       --output-folder datasets/emotion \
+       --n-tasks 16 \
+       hf \
+       --dataset dair-ai/emotion \
 </pre>
 
-The preprocessing script has to be launched with `torchrun` in order to spawn `--nproc-per-node` workers that will preprocess the dataset concurrently. The `--input` dataset can be either a Hugging Face Dataset from the Hub or a `.json` file. The processed dataset will be stored in *`--output-prefix`_input_ids.npy*. In `--tokenizer-name-or-path`, we will have to specify a tokenizer in the same way as we do when using `AutoTokenizers.from_pretrained(...)`.
+First with `--tokenizer-name-or-path` we will specify a tokenizer in the same way as we do when using `AutoTokenizers.from_pretrained(...)`. Then we specify the `--output-folder` where we will store the tokenized documents and the number of workers with `--n-tasks`. Finally we will indicate the type of dataset (whether if it's a Hugging Face Dataset ["**hf**"] or in jsonl ["**jsonl**"] format) and the dataset that we want to preprocess. Check the different settings with `python3 tools/preprocess_data.py --help`, `python3 tools/preprocess_data.py hf --help` & `python3 tools/preprocess_data.py jsonl --help`.
 
-The output will be one file named, in this case, `datasets/testing_alpaca_small_input_ids.npy`. We will then have to specify this file in the `dataset_path` field in the config file.
+Every worker will store in `--output-folder` 3 different kind of files:
+- `*.ds` Containing the tokenized documents
+- `*.ds.index` Containing the bounds of each tokenized document
+- `*.ds.metadata` Containing the number of tokens and tokenizer used
+
+> [!IMPORTANT]
+Remember to introduce the type of dataset to process. e.g. python3 tools/preprocess_data.py --tokenizer-name-or-path gpt2 --n-tasks 16 **jsonl** --dataset raw_datasets/c4-es-json-files
 
 ## Working with Nanosets
 
 To work with `Nanosets`, we just need to configure 1 argument:
-1. `dataset_path`: This argument specifies the file or files that will compose the `Nanoset`. There are 3 ways to specify it:
+1. `dataset_folder`: This argument specifies the file or files that will compose the `Nanoset`. There are 3 ways to specify it:
    1. If we specify a single path, we will create a `Nanoset` from a single dataset file.
     ```yaml
     data_stages:
@@ -43,7 +44,7 @@ To work with `Nanosets`, we just need to configure 1 argument:
         start_training_step: 1
         data:
           dataset:
-            dataset_path: datasets/SlimPajama-6B_input_ids.npy
+            dataset_folder: datasets/SlimPajama-6B
           num_loading_workers: 0
           seed: 1234
     ```
@@ -54,9 +55,9 @@ To work with `Nanosets`, we just need to configure 1 argument:
         start_training_step: 15
         data:
           dataset:
-            dataset_path:
-            - datasets/SlimPajama-6B_input_ids.npy
-            - datasets/testing_alpaca_small_input_ids.npy
+            dataset_folder:
+            - datasets/SlimPajama-6B
+            - datasets/testing_alpaca_small
           num_loading_workers: 0
           seed: 1234
     ```
@@ -67,9 +68,9 @@ To work with `Nanosets`, we just need to configure 1 argument:
         start_training_step: 25
         data:
           dataset:
-            dataset_path:
-              datasets/SlimPajama-6B_input_ids.npy: 0.8
-              datasets/testing_alpaca_small_input_ids.npy: 0.2
+            dataset_folder:
+              datasets/SlimPajama-6B: 0.8
+              datasets/testing_alpaca_small: 0.2
           num_loading_workers: 0
           seed: 1234
     ```
@@ -82,7 +83,10 @@ torchrun --nproc-per-node 8 run_train.py --config configs/config_nanoset.yaml
 ```
 
 ## Under the hood
-`Nanosets` are responsible of building samples of `sequence length + 1` tokens from the preprocessed dataset files. The `dataset lengths` of each dataset will be determined by the `(dataset_number_of_tokens - 1) / sequence length`, discarding the last sample if its length < `sequence length`.
+`Nanosets` are responsible of building samples of `sequence length + 1` tokens from the preprocessed dataset files. Despite most of the extracting logic lies in `DatatroveFolderDataset`, `Nanosets` will take care of the following:
+1. Creating dataset mixtures from different dataset folder paths
+2. Ensure that in each epoch, we consume each sample only once
+3. Ensure that we never exhaust the `DataLoader`
 
 Based on the `dataset lengths`, the `dataset weights` and the `number of samples per epoch` (defined as the `sum(dataset lengths)`), we build the two indexes we need in order to extract samples from the `Nanoset`  ([build_nanoset_index_helper](../src/nanotron/data/nanoset.py)):
 - `dataset index`: Contains the index of the dataset from the list of `dataset paths` from which to extract the sample, respecting the established dataset weight.

From d7cfc3f2b14bc35a092bea3e6c9f605f729616e5 Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Sun, 2 Jun 2024 18:33:42 +0000
Subject: [PATCH 09/33] Install datatrove from source

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 898d22bf..6a0cfb83 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,7 +49,7 @@ fast-modeling = [
 
 nanosets = [
      "transformers",
-     "datatrove[io,processing]",
+     "datatrove[io,processing]@git+https://github.com/huggingface/datatrove",
      "numba",
 ]
 

From bcf405d9af2028773d6d76cd4ff658540b87a3f1 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Thu, 27 Jun 2024 11:56:53 +0000
Subject: [PATCH 10/33] Implemented global memory buffer to reduce activation
 memory of differentiable distributed operations

---
 src/nanotron/config/config.py                 |  2 +-
 .../distributed_differentiable_primitives.py  | 17 +++--------------
 src/nanotron/parallel/utils.py                | 19 +++++++++++++++++++
 src/nanotron/utils.py                         | 17 +++++++++++++++++
 4 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py
index d5b9976f..d72ea97f 100644
--- a/src/nanotron/config/config.py
+++ b/src/nanotron/config/config.py
@@ -111,7 +111,7 @@ def __post_init__(self):
 class DataArgs:
     """Arguments related to the data and data files processing"""
 
-    dataset: Union[PretrainDatasetsArgs, NanosetDatasetsArgs]
+    dataset: Optional[Union[PretrainDatasetsArgs, NanosetDatasetsArgs]]
     seed: Optional[int]
     num_loading_workers: Optional[int] = 1
 
diff --git a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
index 873d77df..57a67c42 100644
--- a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
+++ b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
@@ -19,6 +19,7 @@
 
 from nanotron import distributed as dist
 from nanotron.distributed import ProcessGroup
+from nanotron.parallel.utils import MemoryBuffer
 
 
 class DifferentiableIdentity(torch.autograd.Function):
@@ -67,13 +68,7 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
             group = torch_dist.distributed_c10d._get_default_group()
         unsharded_batch_size = sharded_batch_size * group.size()
 
-        unsharded_tensor = torch.empty(
-            unsharded_batch_size,
-            *rest_size,
-            device=tensor.device,
-            dtype=tensor.dtype,
-            requires_grad=tensor.requires_grad,
-        )
+        unsharded_tensor = MemoryBuffer().get("dist", (unsharded_batch_size, *rest_size), dtype=tensor.dtype)
 
         # `tensor` can sometimes not be contiguous
         # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L317
@@ -108,13 +103,7 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
         # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L305
         tensor = tensor.contiguous()
 
-        sharded_tensor = torch.empty(
-            unsharded_batch_size // group.size(),
-            *rest_size,
-            device=tensor.device,
-            dtype=tensor.dtype,
-            requires_grad=tensor.requires_grad,
-        )
+        sharded_tensor = MemoryBuffer().get("dist", (unsharded_batch_size//group.size(), *rest_size), dtype=tensor.dtype)
         dist.reduce_scatter_tensor(sharded_tensor, tensor, group=group, op=dist.ReduceOp.SUM)
         return sharded_tensor
 
diff --git a/src/nanotron/parallel/utils.py b/src/nanotron/parallel/utils.py
index b9ac12ae..eb4e441d 100644
--- a/src/nanotron/parallel/utils.py
+++ b/src/nanotron/parallel/utils.py
@@ -1,13 +1,32 @@
 import functools
+import operator
 import os
 
+import torch
 from torch import nn
 
 from nanotron import distributed as dist
+from nanotron.utils import Singleton
 from nanotron.parallel import ParallelContext
 from nanotron.parallel.tied_parameters import get_tied_id_to_param
 
 
+class MemoryBuffer(metaclass=Singleton):
+    """
+    Global memory buffer to store intermediate activations that need not to be cached for the backward pass.
+    """
+
+    def __init__(self):
+        self.buffer = {}
+
+    def get(self, name: str, shape: tuple[int], dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        required_numel = functools.reduce(operator.mul, shape, 1)
+        if (name, dtype) not in self.buffer or self.buffer[name, dtype].numel() < required_numel:
+            self.buffer[name, dtype] = torch.empty(required_numel, dtype=dtype, device=torch.cuda.current_device(),
+                                                   requires_grad=False)
+        return self.buffer[name, dtype][:required_numel].view(shape)
+
+
 def assert_cuda_max_connections_set_to_1(func):
     flag_is_set_to_1 = None
 
diff --git a/src/nanotron/utils.py b/src/nanotron/utils.py
index 14fe1ca8..8065962b 100644
--- a/src/nanotron/utils.py
+++ b/src/nanotron/utils.py
@@ -15,6 +15,23 @@
 from nanotron import distributed as dist
 
 
+class Singleton(type):
+    """
+    Singleton metaclass.
+    Create objects using this class as the metaclass to enable singleton behaviour.
+    For instance:
+    ```
+    class Logger(metaclass=Singleton):
+      ...
+    ```
+    """
+    _instances = {}
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+
 class ContextManagers:
     """
     Wrapper for `contextlib.ExitStack` which enters a collection of context managers. Adaptation of `ContextManagers`

From ed1ca7d0b55b07696d1c622d713c793f3ca53e28 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Thu, 27 Jun 2024 14:29:14 +0000
Subject: [PATCH 11/33] GLU fusion

---
 src/nanotron/models/llama.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py
index ca8894b9..d310fe2a 100644
--- a/src/nanotron/models/llama.py
+++ b/src/nanotron/models/llama.py
@@ -163,8 +163,7 @@ def __init__(
             bias=False,
             async_communication=tp_linear_async_communication and tp_mode is TensorParallelLinearMode.REDUCE_SCATTER,
         )
-        # TODO @nouamane: why can't we torch.jit.script GLUActivation?
-        self.split_silu_mul = GLUActivation(config.hidden_act)
+        self.split_silu_mul = torch.compile(GLUActivation(config.hidden_act))
 
     def forward(self, hidden_states):  # [seq_length, batch_size, hidden_dim]
         merged_states = self.gate_up_proj(hidden_states)

From 9b0de5be04afb9cac631399593aef8de6aa852a6 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Thu, 27 Jun 2024 14:42:42 +0000
Subject: [PATCH 12/33] precommit

---
 src/nanotron/models/llama.py                              | 2 +-
 .../distributed_differentiable_primitives.py              | 4 +++-
 src/nanotron/parallel/utils.py                            | 7 ++++---
 src/nanotron/utils.py                                     | 8 +++++---
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py
index d310fe2a..3319b0ef 100644
--- a/src/nanotron/models/llama.py
+++ b/src/nanotron/models/llama.py
@@ -14,7 +14,7 @@
 # limitations under the License.
 """PyTorch LLaMa model."""
 
-from typing import Dict, Optional, Union, List
+from typing import Dict, Optional, Union
 
 import torch
 from torch import nn
diff --git a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
index 57a67c42..aa460cc6 100644
--- a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
+++ b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
@@ -103,7 +103,9 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
         # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L305
         tensor = tensor.contiguous()
 
-        sharded_tensor = MemoryBuffer().get("dist", (unsharded_batch_size//group.size(), *rest_size), dtype=tensor.dtype)
+        sharded_tensor = MemoryBuffer().get(
+            "dist", (unsharded_batch_size // group.size(), *rest_size), dtype=tensor.dtype
+        )
         dist.reduce_scatter_tensor(sharded_tensor, tensor, group=group, op=dist.ReduceOp.SUM)
         return sharded_tensor
 
diff --git a/src/nanotron/parallel/utils.py b/src/nanotron/parallel/utils.py
index eb4e441d..f694b0e6 100644
--- a/src/nanotron/parallel/utils.py
+++ b/src/nanotron/parallel/utils.py
@@ -6,9 +6,9 @@
 from torch import nn
 
 from nanotron import distributed as dist
-from nanotron.utils import Singleton
 from nanotron.parallel import ParallelContext
 from nanotron.parallel.tied_parameters import get_tied_id_to_param
+from nanotron.utils import Singleton
 
 
 class MemoryBuffer(metaclass=Singleton):
@@ -22,8 +22,9 @@ def __init__(self):
     def get(self, name: str, shape: tuple[int], dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
         required_numel = functools.reduce(operator.mul, shape, 1)
         if (name, dtype) not in self.buffer or self.buffer[name, dtype].numel() < required_numel:
-            self.buffer[name, dtype] = torch.empty(required_numel, dtype=dtype, device=torch.cuda.current_device(),
-                                                   requires_grad=False)
+            self.buffer[name, dtype] = torch.empty(
+                required_numel, dtype=dtype, device=torch.cuda.current_device(), requires_grad=False
+            )
         return self.buffer[name, dtype][:required_numel].view(shape)
 
 
diff --git a/src/nanotron/utils.py b/src/nanotron/utils.py
index 8065962b..b3831801 100644
--- a/src/nanotron/utils.py
+++ b/src/nanotron/utils.py
@@ -1,11 +1,10 @@
 import functools
 import inspect
-import math
 import os
 import random
 import socket
 from contextlib import ExitStack, contextmanager
-from typing import Callable, ContextManager, List, Optional
+from typing import ContextManager, List, Optional
 
 import torch
 from packaging import version
@@ -25,7 +24,9 @@ class Logger(metaclass=Singleton):
       ...
     ```
     """
+
     _instances = {}
+
     def __call__(cls, *args, **kwargs):
         if cls not in cls._instances:
             cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
@@ -69,7 +70,7 @@ def main_rank_first(group: dist.ProcessGroup):
 @contextmanager
 def local_ranks_zero_first(group: Optional[dist.ProcessGroup] = None):
     """Context manager that executes the code in the context with all the local rank zero of the group going first.
-    Usefull to run only once per node first (e.g. to create local files, etc)
+    Useful to run only once per node first (e.g. to create local files, etc)
     """
     is_main = int(os.environ.get("LOCAL_RANK", 0)) == 0
     if is_main:
@@ -140,6 +141,7 @@ def get_untyped_storage(tensor: torch.Tensor) -> torch.UntypedStorage:
     else:
         return tensor.storage().untyped()
 
+
 def tensor_from_untyped_storage(untyped_storage: torch.UntypedStorage, dtype: torch.dtype):
     # TODO @thomasw21: Figure out what's the best Pytorch way of building a tensor from a storage.
     device = untyped_storage.device

From 803b6da3233a642a0ba7a62484310d1496db81dc Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 16 Jul 2024 11:39:32 +0200
Subject: [PATCH 13/33] Wrong backward fixed

---
 .../parallel/tensor_parallel/column_linear.py | 62 +++++++++++++++++++
 .../distributed_differentiable_primitives.py  | 27 +++++---
 .../parallel/tensor_parallel/functional.py    |  4 +-
 3 files changed, 85 insertions(+), 8 deletions(-)
 create mode 100644 src/nanotron/parallel/tensor_parallel/column_linear.py

diff --git a/src/nanotron/parallel/tensor_parallel/column_linear.py b/src/nanotron/parallel/tensor_parallel/column_linear.py
new file mode 100644
index 00000000..eaab5abe
--- /dev/null
+++ b/src/nanotron/parallel/tensor_parallel/column_linear.py
@@ -0,0 +1,62 @@
+from typing import Optional
+
+import torch
+from torch.nn import functional as F
+
+import nanotron.distributed as dist
+from nanotron.parallel.utils import MemoryBuffer
+
+
+class ColumnLinearContextParallel(torch.autograd.Function):
+    """
+    Column linear with memory_buffer for the allgather, context parallel
+    enabled (i.e. tp_mode = TensorParallelLinearMode.REDUCE_SCATTER) and
+    async communication disabled.
+    """
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight: torch.Tensor,
+                bias: Optional[torch.Tensor], group: dist.ProcessGroup):
+
+        # Prepare context.
+        ctx.save_for_backward(input, weight, bias)
+        ctx.group = group
+
+        # Do allgather.
+        sharded_batch_size, *rest_size = input.shape
+        unsharded_batch_size = sharded_batch_size * group.size()
+        total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+        dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+
+        # Get linear output.
+        out = F.linear(total_input, weight, bias)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        # Allgather the inputs again.
+        input, weight, bias = ctx.saved_tensors
+        group = ctx.group
+        sharded_batch_size, *rest_size = input.shape
+        total_input = sharded_batch_size * group.size()
+        unsharded_batch_size = sharded_batch_size * group.size()
+        total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+        dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+
+        # Get the grad_output and total_input on the correct views to be able to transpose them below.
+        grad_output = grad_output.contiguous()
+        assert grad_output.dim() == 3
+        grad_output = grad_output.view(grad_output.size(0) * grad_output.size(1), grad_output.size(2))
+        total_input = total_input.view(total_input.size(0) * total_input.size(1), total_input.size(2))
+
+        # Compute gradients.
+        grad_input = grad_output @ weight
+        sub_grad_input = torch.empty(input.size(), dtype=input.dtype, device=input.device, requires_grad=False)
+        dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
+        grad_weight = grad_output.T @ total_input
+        grad_bias = torch.sum(grad_output, dim=0) if bias is not None else None
+
+        return sub_grad_input, grad_weight, grad_bias, None
+
+def column_linear_context_parallel(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                                   group: dist.ProcessGroup):
+    return ColumnLinearContextParallel.apply(input, weight, bias, group)
diff --git a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
index aa460cc6..d66826e3 100644
--- a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
+++ b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
@@ -19,7 +19,6 @@
 
 from nanotron import distributed as dist
 from nanotron.distributed import ProcessGroup
-from nanotron.parallel.utils import MemoryBuffer
 
 
 class DifferentiableIdentity(torch.autograd.Function):
@@ -68,7 +67,13 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
             group = torch_dist.distributed_c10d._get_default_group()
         unsharded_batch_size = sharded_batch_size * group.size()
 
-        unsharded_tensor = MemoryBuffer().get("dist", (unsharded_batch_size, *rest_size), dtype=tensor.dtype)
+        unsharded_tensor = torch.empty(
+            unsharded_batch_size,
+            *rest_size,
+            device=tensor.device,
+            dtype=tensor.dtype,
+            requires_grad=tensor.requires_grad,
+        )
 
         # `tensor` can sometimes not be contiguous
         # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L317
@@ -79,8 +84,11 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
 
     @staticmethod
     def backward(ctx, grad_output):
+        #print(f"{torch.distributed.get_rank()} grad_output: {grad_output}")
         group = ctx.group
-        return DifferentiableReduceScatterSum.apply(grad_output, group), None
+        out = DifferentiableReduceScatterSum.apply(grad_output, group)
+        #print(f"{torch.distributed.get_rank()} grad_grad: {out}")
+        return out, None, None
 
 
 class DifferentiableReduceScatterSum(torch.autograd.Function):
@@ -103,8 +111,12 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
         # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L305
         tensor = tensor.contiguous()
 
-        sharded_tensor = MemoryBuffer().get(
-            "dist", (unsharded_batch_size // group.size(), *rest_size), dtype=tensor.dtype
+        sharded_tensor = torch.empty(
+            unsharded_batch_size // group.size(),
+            *rest_size,
+            device=tensor.device,
+            dtype=tensor.dtype,
+            requires_grad=False,
         )
         dist.reduce_scatter_tensor(sharded_tensor, tensor, group=group, op=dist.ReduceOp.SUM)
         return sharded_tensor
@@ -112,7 +124,8 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
     @staticmethod
     def backward(ctx, grad_output):
         group = ctx.group
-        return DifferentiableAllGather.apply(grad_output, group), None
+        #print(f"{torch.distributed.get_rank()} Calling AllGather because of backward of reducescatter")
+        return DifferentiableAllGather.apply(grad_output, group, False), None
 
 
 # -----------------
@@ -128,7 +141,7 @@ def differentiable_all_reduce_sum(tensor, group: Optional[ProcessGroup] = None):
     return DifferentiableAllReduceSum.apply(tensor, group)
 
 
-def differentiable_all_gather(tensor, group: Optional[ProcessGroup] = None):
+def differentiable_all_gather(tensor, group: Optional[ProcessGroup] = None)
     return DifferentiableAllGather.apply(tensor, group)
 
 
diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index fdef48ac..b3602707 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -26,6 +26,7 @@
     differentiable_reduce_scatter_sum,
 )
 from nanotron.parallel.tensor_parallel.enum import TensorParallelLinearMode
+from nanotron.parallel.tensor_parallel.column_linear import column_linear_context_parallel
 from nanotron.parallel.utils import assert_cuda_max_connections_set_to_1
 
 
@@ -352,7 +353,7 @@ def column_linear(
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         input = differentiable_identity(input, group=group)
     elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
-        input = differentiable_all_gather(input, group=group)
+        return column_linear_context_parallel(input, weight, bias, group)
     else:
         raise ValueError(f"Got unexpected mode: {tp_mode}.")
 
@@ -473,6 +474,7 @@ def row_linear(
 
     out = F.linear(input, weight, bias)
 
+    #print("Calling row linear")
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         out = differentiable_all_reduce_sum(out, group=group)
     elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:

From 59bfb6b38c4487b04e425d8540b6e44b2a7fbcf9 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 16 Jul 2024 11:42:27 +0200
Subject: [PATCH 14/33] Removed useless prints

---
 .../tensor_parallel/distributed_differentiable_primitives.py   | 3 ---
 src/nanotron/parallel/tensor_parallel/functional.py            | 1 -
 2 files changed, 4 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
index d66826e3..f1102908 100644
--- a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
+++ b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
@@ -84,10 +84,8 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
 
     @staticmethod
     def backward(ctx, grad_output):
-        #print(f"{torch.distributed.get_rank()} grad_output: {grad_output}")
         group = ctx.group
         out = DifferentiableReduceScatterSum.apply(grad_output, group)
-        #print(f"{torch.distributed.get_rank()} grad_grad: {out}")
         return out, None, None
 
 
@@ -124,7 +122,6 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
     @staticmethod
     def backward(ctx, grad_output):
         group = ctx.group
-        #print(f"{torch.distributed.get_rank()} Calling AllGather because of backward of reducescatter")
         return DifferentiableAllGather.apply(grad_output, group, False), None
 
 
diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index b3602707..cedbb219 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -474,7 +474,6 @@ def row_linear(
 
     out = F.linear(input, weight, bias)
 
-    #print("Calling row linear")
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         out = differentiable_all_reduce_sum(out, group=group)
     elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:

From 2c69e9ad2887b7e78c88c2db3209713542dad7e2 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Wed, 17 Jul 2024 10:01:44 +0000
Subject: [PATCH 15/33] Minor fixes

---
 .../distributed_differentiable_primitives.py                | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
index f1102908..bd41347a 100644
--- a/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
+++ b/src/nanotron/parallel/tensor_parallel/distributed_differentiable_primitives.py
@@ -86,7 +86,7 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
     def backward(ctx, grad_output):
         group = ctx.group
         out = DifferentiableReduceScatterSum.apply(grad_output, group)
-        return out, None, None
+        return out, None
 
 
 class DifferentiableReduceScatterSum(torch.autograd.Function):
@@ -122,7 +122,7 @@ def forward(ctx, tensor, group: Optional[ProcessGroup]):
     @staticmethod
     def backward(ctx, grad_output):
         group = ctx.group
-        return DifferentiableAllGather.apply(grad_output, group, False), None
+        return DifferentiableAllGather.apply(grad_output, group), None
 
 
 # -----------------
@@ -138,7 +138,7 @@ def differentiable_all_reduce_sum(tensor, group: Optional[ProcessGroup] = None):
     return DifferentiableAllReduceSum.apply(tensor, group)
 
 
-def differentiable_all_gather(tensor, group: Optional[ProcessGroup] = None)
+def differentiable_all_gather(tensor, group: Optional[ProcessGroup] = None):
     return DifferentiableAllGather.apply(tensor, group)
 
 

From 30439fdee7cac456be4a2c28798b42c931f7cf72 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Wed, 17 Jul 2024 11:16:40 +0000
Subject: [PATCH 16/33] precommit

---
 .../parallel/tensor_parallel/column_linear.py        | 12 ++++++++----
 src/nanotron/parallel/tensor_parallel/functional.py  |  7 +++----
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/column_linear.py b/src/nanotron/parallel/tensor_parallel/column_linear.py
index eaab5abe..21daba36 100644
--- a/src/nanotron/parallel/tensor_parallel/column_linear.py
+++ b/src/nanotron/parallel/tensor_parallel/column_linear.py
@@ -13,9 +13,11 @@ class ColumnLinearContextParallel(torch.autograd.Function):
     enabled (i.e. tp_mode = TensorParallelLinearMode.REDUCE_SCATTER) and
     async communication disabled.
     """
+
     @staticmethod
-    def forward(ctx, input: torch.Tensor, weight: torch.Tensor,
-                bias: Optional[torch.Tensor], group: dist.ProcessGroup):
+    def forward(
+        ctx, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup
+    ):
 
         # Prepare context.
         ctx.save_for_backward(input, weight, bias)
@@ -57,6 +59,8 @@ def backward(ctx, grad_output: torch.Tensor):
 
         return sub_grad_input, grad_weight, grad_bias, None
 
-def column_linear_context_parallel(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
-                                   group: dist.ProcessGroup):
+
+def column_linear_context_parallel(
+    input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup
+):
     return ColumnLinearContextParallel.apply(input, weight, bias, group)
diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index cedbb219..f4e9de30 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -19,14 +19,13 @@
 from torch.nn import functional as F
 
 import nanotron.distributed as dist
+from nanotron.parallel.tensor_parallel.column_linear import column_linear_context_parallel
 from nanotron.parallel.tensor_parallel.distributed_differentiable_primitives import (
-    differentiable_all_gather,
     differentiable_all_reduce_sum,
     differentiable_identity,
     differentiable_reduce_scatter_sum,
 )
 from nanotron.parallel.tensor_parallel.enum import TensorParallelLinearMode
-from nanotron.parallel.tensor_parallel.column_linear import column_linear_context_parallel
 from nanotron.parallel.utils import assert_cuda_max_connections_set_to_1
 
 
@@ -90,10 +89,10 @@ def forward(
 
     @staticmethod
     def backward(ctx, grad_output):
-        # Retreive tensors from the forward path.
+        # Retrieve tensors from the forward path.
         softmax, target_mask, masked_target_1d = ctx.saved_tensors
 
-        # All the inputs have softmax as thier gradient.
+        # All the inputs have softmax as their gradient.
         grad_input = softmax
         # For simplicity, work with the 2D gradient.
         sharded_hidden_size = softmax.size()[-1]

From 1e02a9ce9c9b564f4a4274ee62e7208e3d5d9df8 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Wed, 17 Jul 2024 12:47:20 +0000
Subject: [PATCH 17/33] Added tp_recompute_allgather option

---
 src/nanotron/config/parallelism_config.py     |  2 +
 src/nanotron/models/llama.py                  |  3 ++
 .../parallel/tensor_parallel/column_linear.py | 51 ++++++++++++-------
 .../parallel/tensor_parallel/functional.py    |  4 +-
 src/nanotron/parallel/tensor_parallel/nn.py   |  3 ++
 5 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/src/nanotron/config/parallelism_config.py b/src/nanotron/config/parallelism_config.py
index 5912425b..e9a6f2a4 100644
--- a/src/nanotron/config/parallelism_config.py
+++ b/src/nanotron/config/parallelism_config.py
@@ -32,6 +32,8 @@ class ParallelismArgs:
     tp_mode: Optional[TensorParallelLinearMode] = None
     tp_linear_async_communication: Optional[bool] = None
 
+    tp_recompute_allgather: bool = False
+
     expert_parallel_size: int = 1
 
     def __post_init__(self):
diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py
index 3319b0ef..a31ebec6 100644
--- a/src/nanotron/models/llama.py
+++ b/src/nanotron/models/llama.py
@@ -154,6 +154,7 @@ def __init__(
             bias=False,
             async_communication=tp_linear_async_communication,
             contiguous_chunks=gate_up_contiguous_chunks,
+            tp_recompute_allgather=parallel_config.tp_recompute_allgather,
         )
         self.down_proj = TensorParallelRowLinear(
             config.intermediate_size,
@@ -314,6 +315,7 @@ def __init__(
             bias=False,
             async_communication=tp_linear_async_communication,
             contiguous_chunks=qkv_contiguous_chunks,
+            tp_recompute_allgather=parallel_config.tp_recompute_allgather,
         )
         # TODO(kunhao): We want to have only one version per device and not one version per layer.
         self.rotary_embedding = RotaryEmbedding(
@@ -742,6 +744,7 @@ def __init__(
                 # TODO @thomasw21: refactor so that we store that default in a single place.
                 "mode": self.tp_mode,
                 "async_communication": tp_linear_async_communication,
+                "tp_recompute_allgather": parallel_config.tp_recompute_allgather,
             },
             module_input_keys={"x"},
             module_output_keys={"logits"},
diff --git a/src/nanotron/parallel/tensor_parallel/column_linear.py b/src/nanotron/parallel/tensor_parallel/column_linear.py
index 21daba36..2f743199 100644
--- a/src/nanotron/parallel/tensor_parallel/column_linear.py
+++ b/src/nanotron/parallel/tensor_parallel/column_linear.py
@@ -16,33 +16,47 @@ class ColumnLinearContextParallel(torch.autograd.Function):
 
     @staticmethod
     def forward(
-        ctx, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup
+        ctx, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup,
+        tp_recompute_allgather: bool
     ):
 
-        # Prepare context.
-        ctx.save_for_backward(input, weight, bias)
-        ctx.group = group
-
         # Do allgather.
         sharded_batch_size, *rest_size = input.shape
         unsharded_batch_size = sharded_batch_size * group.size()
-        total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+        if tp_recompute_allgather:
+            total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+        else:
+            total_input = torch.empty(unsharded_batch_size, *rest_size, dtype=input.dtype, device=input.device)
         dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
 
+        # Prepare context.
+        ctx.group = group
+        ctx.tp_recompute_allgather = tp_recompute_allgather
+        ctx.input_size = input.shape
+        if tp_recompute_allgather:
+            ctx.save_for_backward(input, weight, bias)
+        else:
+            ctx.save_for_backward(total_input, weight, bias)
+
         # Get linear output.
         out = F.linear(total_input, weight, bias)
         return out
 
     @staticmethod
     def backward(ctx, grad_output: torch.Tensor):
-        # Allgather the inputs again.
-        input, weight, bias = ctx.saved_tensors
+        # Either allgather the inputs again or get them from context.
         group = ctx.group
-        sharded_batch_size, *rest_size = input.shape
-        total_input = sharded_batch_size * group.size()
-        unsharded_batch_size = sharded_batch_size * group.size()
-        total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
-        dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+        tp_recompute_allgather = ctx.tp_recompute_allgather
+        input_size = ctx.input_size
+        if tp_recompute_allgather:
+            input, weight, bias = ctx.saved_tensors
+            sharded_batch_size, *rest_size = input.shape
+            total_input = sharded_batch_size * group.size()
+            unsharded_batch_size = sharded_batch_size * group.size()
+            total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+            dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+        else:
+            total_input, weight, bias = ctx.saved_tensors
 
         # Get the grad_output and total_input on the correct views to be able to transpose them below.
         grad_output = grad_output.contiguous()
@@ -51,16 +65,17 @@ def backward(ctx, grad_output: torch.Tensor):
         total_input = total_input.view(total_input.size(0) * total_input.size(1), total_input.size(2))
 
         # Compute gradients.
+        grad_weight = grad_output.T @ total_input
         grad_input = grad_output @ weight
-        sub_grad_input = torch.empty(input.size(), dtype=input.dtype, device=input.device, requires_grad=False)
+        sub_grad_input = torch.empty(input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False)
         dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
-        grad_weight = grad_output.T @ total_input
         grad_bias = torch.sum(grad_output, dim=0) if bias is not None else None
 
-        return sub_grad_input, grad_weight, grad_bias, None
+        return sub_grad_input, grad_weight, grad_bias, None, None
 
 
 def column_linear_context_parallel(
-    input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup
+    input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup,
+    tp_recompute_allgather: bool = False
 ):
-    return ColumnLinearContextParallel.apply(input, weight, bias, group)
+    return ColumnLinearContextParallel.apply(input, weight, bias, group, tp_recompute_allgather)
diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index f4e9de30..c16ae492 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -21,6 +21,7 @@
 import nanotron.distributed as dist
 from nanotron.parallel.tensor_parallel.column_linear import column_linear_context_parallel
 from nanotron.parallel.tensor_parallel.distributed_differentiable_primitives import (
+    differentiable_all_gather,
     differentiable_all_reduce_sum,
     differentiable_identity,
     differentiable_reduce_scatter_sum,
@@ -345,6 +346,7 @@ def column_linear(
     group: dist.ProcessGroup,
     tp_mode: TensorParallelLinearMode,
     async_communication: bool,
+    tp_recompute_allgather: bool = True
 ):
     if async_communication:
         return _ColumnLinearAsyncCommunication.apply(input, weight, bias, group, tp_mode)
@@ -352,7 +354,7 @@ def column_linear(
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         input = differentiable_identity(input, group=group)
     elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
-        return column_linear_context_parallel(input, weight, bias, group)
+        return column_linear_context_parallel(input, weight, bias, group, tp_recompute_allgather)
     else:
         raise ValueError(f"Got unexpected mode: {tp_mode}.")
 
diff --git a/src/nanotron/parallel/tensor_parallel/nn.py b/src/nanotron/parallel/tensor_parallel/nn.py
index 40e89968..42ffc828 100644
--- a/src/nanotron/parallel/tensor_parallel/nn.py
+++ b/src/nanotron/parallel/tensor_parallel/nn.py
@@ -51,6 +51,7 @@ def __init__(
         dtype=None,
         async_communication: bool = False,
         contiguous_chunks: Optional[Tuple[int, ...]] = None,
+        tp_recompute_allgather: bool = False,
     ):
         self.pg = pg
         self.world_size = pg.size()
@@ -59,6 +60,7 @@ def __init__(
 
         self.in_features = in_features
         self.out_features = out_features // self.world_size
+        self.tp_recompute_allgather = tp_recompute_allgather
 
         super().__init__(
             in_features=self.in_features,
@@ -91,6 +93,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             group=self.pg,
             tp_mode=self.mode,
             async_communication=self.async_communication,
+            tp_recompute_allgather=self.tp_recompute_allgather,
         )
 
     def extra_repr(self) -> str:

From 9cc81bb6fe680b72cf6114f7258af0483886ada1 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Wed, 17 Jul 2024 13:39:14 +0000
Subject: [PATCH 18/33] Changed recompute default

---
 src/nanotron/config/parallelism_config.py     |  2 +-
 .../parallel/tensor_parallel/column_linear.py | 19 ++++++++++++++-----
 .../parallel/tensor_parallel/functional.py    |  3 +--
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/nanotron/config/parallelism_config.py b/src/nanotron/config/parallelism_config.py
index e9a6f2a4..cc5d406a 100644
--- a/src/nanotron/config/parallelism_config.py
+++ b/src/nanotron/config/parallelism_config.py
@@ -32,7 +32,7 @@ class ParallelismArgs:
     tp_mode: Optional[TensorParallelLinearMode] = None
     tp_linear_async_communication: Optional[bool] = None
 
-    tp_recompute_allgather: bool = False
+    tp_recompute_allgather: bool = True
 
     expert_parallel_size: int = 1
 
diff --git a/src/nanotron/parallel/tensor_parallel/column_linear.py b/src/nanotron/parallel/tensor_parallel/column_linear.py
index 2f743199..880d5ff0 100644
--- a/src/nanotron/parallel/tensor_parallel/column_linear.py
+++ b/src/nanotron/parallel/tensor_parallel/column_linear.py
@@ -16,8 +16,12 @@ class ColumnLinearContextParallel(torch.autograd.Function):
 
     @staticmethod
     def forward(
-        ctx, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup,
-        tp_recompute_allgather: bool
+        ctx,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        group: dist.ProcessGroup,
+        tp_recompute_allgather: bool,
     ):
 
         # Do allgather.
@@ -67,7 +71,9 @@ def backward(ctx, grad_output: torch.Tensor):
         # Compute gradients.
         grad_weight = grad_output.T @ total_input
         grad_input = grad_output @ weight
-        sub_grad_input = torch.empty(input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False)
+        sub_grad_input = torch.empty(
+            input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False
+        )
         dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
         grad_bias = torch.sum(grad_output, dim=0) if bias is not None else None
 
@@ -75,7 +81,10 @@ def backward(ctx, grad_output: torch.Tensor):
 
 
 def column_linear_context_parallel(
-    input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], group: dist.ProcessGroup,
-    tp_recompute_allgather: bool = False
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    group: dist.ProcessGroup,
+    tp_recompute_allgather: bool = True,
 ):
     return ColumnLinearContextParallel.apply(input, weight, bias, group, tp_recompute_allgather)
diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index c16ae492..454cc447 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -21,7 +21,6 @@
 import nanotron.distributed as dist
 from nanotron.parallel.tensor_parallel.column_linear import column_linear_context_parallel
 from nanotron.parallel.tensor_parallel.distributed_differentiable_primitives import (
-    differentiable_all_gather,
     differentiable_all_reduce_sum,
     differentiable_identity,
     differentiable_reduce_scatter_sum,
@@ -346,7 +345,7 @@ def column_linear(
     group: dist.ProcessGroup,
     tp_mode: TensorParallelLinearMode,
     async_communication: bool,
-    tp_recompute_allgather: bool = True
+    tp_recompute_allgather: bool = True,
 ):
     if async_communication:
         return _ColumnLinearAsyncCommunication.apply(input, weight, bias, group, tp_mode)

From 956fbfd09a2f0c2358fcb90be395f97ffa79632e Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Wed, 17 Jul 2024 13:40:34 +0000
Subject: [PATCH 19/33] Changed recompute default

---
 src/nanotron/parallel/tensor_parallel/nn.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/nanotron/parallel/tensor_parallel/nn.py b/src/nanotron/parallel/tensor_parallel/nn.py
index 42ffc828..4c7325cd 100644
--- a/src/nanotron/parallel/tensor_parallel/nn.py
+++ b/src/nanotron/parallel/tensor_parallel/nn.py
@@ -51,7 +51,7 @@ def __init__(
         dtype=None,
         async_communication: bool = False,
         contiguous_chunks: Optional[Tuple[int, ...]] = None,
-        tp_recompute_allgather: bool = False,
+        tp_recompute_allgather: bool = True,
     ):
         self.pg = pg
         self.world_size = pg.size()

From 9992f1c5919fd4038e85cd9b3fb1dd4faa81daf1 Mon Sep 17 00:00:00 2001
From: tj-solergibert <tj.solergibert@gmail.com>
Date: Wed, 17 Jul 2024 14:28:57 +0000
Subject: [PATCH 20/33] Little fixes

---
 src/nanotron/config/config.py | 4 ++--
 tools/preprocess_data.py      | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py
index fe194883..2e1a98cc 100644
--- a/src/nanotron/config/config.py
+++ b/src/nanotron/config/config.py
@@ -96,10 +96,10 @@ class NanosetDatasetsArgs:
     dataset_folder: Union[str, dict, List[str]]
 
     def __post_init__(self):
-        if isinstance(self.dataset_folder, str):  # Case 1: 1 Dataset file
+        if isinstance(self.dataset_folder, str):  # Case 1: 1 Dataset folder
             self.dataset_folder = [self.dataset_folder]
             self.dataset_weights = [1]
-        elif isinstance(self.dataset_folder, List):  # Case 2: > 1 Dataset file
+        elif isinstance(self.dataset_folder, List):  # Case 2: > 1 Dataset folder
             self.dataset_weights = None  # Set to None so we consume all the samples randomly
         elif isinstance(self.dataset_folder, dict):  # Case 3: dict with > 1 dataset_folder and weights
             tmp_dataset_folder = self.dataset_folder.copy()
diff --git a/tools/preprocess_data.py b/tools/preprocess_data.py
index 38db67f1..f3cdab70 100644
--- a/tools/preprocess_data.py
+++ b/tools/preprocess_data.py
@@ -95,6 +95,7 @@ def main(args):
                 output_folder=args.output_folder,
                 tokenizer_name_or_path=args.tokenizer_name_or_path,
                 eos_token=args.eos_token,
+                shuffle=False,
                 max_tokens_per_file=1e9,
             ),
         ],

From b9e92017614e0326acab86b7665e0c8e8718bfc3 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Wed, 17 Jul 2024 15:21:04 +0000
Subject: [PATCH 21/33] Moved ColumnLinearNoAsync module for consistency

---
 .../parallel/tensor_parallel/column_linear.py | 90 -------------------
 .../parallel/tensor_parallel/functional.py    | 78 +++++++++++++++-
 2 files changed, 76 insertions(+), 92 deletions(-)
 delete mode 100644 src/nanotron/parallel/tensor_parallel/column_linear.py

diff --git a/src/nanotron/parallel/tensor_parallel/column_linear.py b/src/nanotron/parallel/tensor_parallel/column_linear.py
deleted file mode 100644
index 880d5ff0..00000000
--- a/src/nanotron/parallel/tensor_parallel/column_linear.py
+++ /dev/null
@@ -1,90 +0,0 @@
-from typing import Optional
-
-import torch
-from torch.nn import functional as F
-
-import nanotron.distributed as dist
-from nanotron.parallel.utils import MemoryBuffer
-
-
-class ColumnLinearContextParallel(torch.autograd.Function):
-    """
-    Column linear with memory_buffer for the allgather, context parallel
-    enabled (i.e. tp_mode = TensorParallelLinearMode.REDUCE_SCATTER) and
-    async communication disabled.
-    """
-
-    @staticmethod
-    def forward(
-        ctx,
-        input: torch.Tensor,
-        weight: torch.Tensor,
-        bias: Optional[torch.Tensor],
-        group: dist.ProcessGroup,
-        tp_recompute_allgather: bool,
-    ):
-
-        # Do allgather.
-        sharded_batch_size, *rest_size = input.shape
-        unsharded_batch_size = sharded_batch_size * group.size()
-        if tp_recompute_allgather:
-            total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
-        else:
-            total_input = torch.empty(unsharded_batch_size, *rest_size, dtype=input.dtype, device=input.device)
-        dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
-
-        # Prepare context.
-        ctx.group = group
-        ctx.tp_recompute_allgather = tp_recompute_allgather
-        ctx.input_size = input.shape
-        if tp_recompute_allgather:
-            ctx.save_for_backward(input, weight, bias)
-        else:
-            ctx.save_for_backward(total_input, weight, bias)
-
-        # Get linear output.
-        out = F.linear(total_input, weight, bias)
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):
-        # Either allgather the inputs again or get them from context.
-        group = ctx.group
-        tp_recompute_allgather = ctx.tp_recompute_allgather
-        input_size = ctx.input_size
-        if tp_recompute_allgather:
-            input, weight, bias = ctx.saved_tensors
-            sharded_batch_size, *rest_size = input.shape
-            total_input = sharded_batch_size * group.size()
-            unsharded_batch_size = sharded_batch_size * group.size()
-            total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
-            dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
-        else:
-            total_input, weight, bias = ctx.saved_tensors
-
-        # Get the grad_output and total_input on the correct views to be able to transpose them below.
-        grad_output = grad_output.contiguous()
-        assert grad_output.dim() == 3
-        grad_output = grad_output.view(grad_output.size(0) * grad_output.size(1), grad_output.size(2))
-        total_input = total_input.view(total_input.size(0) * total_input.size(1), total_input.size(2))
-
-        # Compute gradients.
-        grad_weight = grad_output.T @ total_input
-        grad_input = grad_output @ weight
-        sub_grad_input = torch.empty(
-            input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False
-        )
-        dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
-        grad_bias = torch.sum(grad_output, dim=0) if bias is not None else None
-
-        return sub_grad_input, grad_weight, grad_bias, None, None
-
-
-def column_linear_context_parallel(
-    input: torch.Tensor,
-    weight: torch.Tensor,
-    bias: Optional[torch.Tensor],
-    group: dist.ProcessGroup,
-    tp_recompute_allgather: bool = True,
-):
-    return ColumnLinearContextParallel.apply(input, weight, bias, group, tp_recompute_allgather)
diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 454cc447..2b93fb02 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -19,7 +19,7 @@
 from torch.nn import functional as F
 
 import nanotron.distributed as dist
-from nanotron.parallel.tensor_parallel.column_linear import column_linear_context_parallel
+from nanotron.parallel.utils import MemoryBuffer
 from nanotron.parallel.tensor_parallel.distributed_differentiable_primitives import (
     differentiable_all_reduce_sum,
     differentiable_identity,
@@ -338,6 +338,80 @@ def backward(ctx, grad_output):
             raise ValueError(f"Got unexpected mode: {tp_mode}.")
 
 
+class _ColumnLinearContextParallelNoAsync(torch.autograd.Function):
+    """
+    Column linear with memory_buffer for the allgather, context parallel
+    enabled (i.e. tp_mode = TensorParallelLinearMode.REDUCE_SCATTER) and
+    async communication disabled.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        group: dist.ProcessGroup,
+        tp_recompute_allgather: bool,
+    ):
+
+        # Do allgather.
+        sharded_batch_size, *rest_size = input.shape
+        unsharded_batch_size = sharded_batch_size * group.size()
+        if tp_recompute_allgather:
+            total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+        else:
+            total_input = torch.empty(unsharded_batch_size, *rest_size, dtype=input.dtype, device=input.device)
+        dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+
+        # Prepare context.
+        ctx.group = group
+        ctx.tp_recompute_allgather = tp_recompute_allgather
+        ctx.input_size = input.shape
+        if tp_recompute_allgather:
+            ctx.save_for_backward(input, weight, bias)
+        else:
+            ctx.save_for_backward(total_input, weight, bias)
+
+        # Get linear output.
+        out = F.linear(total_input, weight, bias)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        # Either allgather the inputs again or get them from context.
+        group = ctx.group
+        tp_recompute_allgather = ctx.tp_recompute_allgather
+        input_size = ctx.input_size
+        if tp_recompute_allgather:
+            input, weight, bias = ctx.saved_tensors
+            sharded_batch_size, *rest_size = input.shape
+            total_input = sharded_batch_size * group.size()
+            unsharded_batch_size = sharded_batch_size * group.size()
+            total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+            dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+        else:
+            total_input, weight, bias = ctx.saved_tensors
+
+        # Get the grad_output and total_input on the correct views to be able to transpose them below.
+        grad_output = grad_output.contiguous()
+        assert grad_output.dim() == 3
+        grad_output = grad_output.view(grad_output.size(0) * grad_output.size(1), grad_output.size(2))
+        total_input = total_input.view(total_input.size(0) * total_input.size(1), total_input.size(2))
+
+        # Compute gradients.
+        grad_weight = grad_output.T @ total_input
+        grad_input = grad_output @ weight
+        sub_grad_input = torch.empty(
+            input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False
+        )
+        dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
+        grad_bias = torch.sum(grad_output, dim=0) if bias is not None else None
+
+        return sub_grad_input, grad_weight, grad_bias, None, None
+
+
+
 def column_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -353,7 +427,7 @@ def column_linear(
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         input = differentiable_identity(input, group=group)
     elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
-        return column_linear_context_parallel(input, weight, bias, group, tp_recompute_allgather)
+        return _ColumnLinearContextParallelNoAsync.apply(input, weight, bias, group, tp_recompute_allgather)
     else:
         raise ValueError(f"Got unexpected mode: {tp_mode}.")
 

From 7cc6653c69c19cd42da05e6a8712159a146407e7 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Thu, 18 Jul 2024 11:16:42 +0000
Subject: [PATCH 22/33] memory efficient async linear

---
 .../parallel/tensor_parallel/functional.py    | 48 ++++++-------------
 1 file changed, 15 insertions(+), 33 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 47c0b5a1..1a82254e 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -149,14 +149,7 @@ def forward(ctx, tensor, weight, bias, group, tp_mode):
                     group = dist.distributed_c10d._get_default_group()
                 gathered_batch_size = sharded_batch_size * group.size()
 
-                gathered_tensor = torch.empty(
-                    gathered_batch_size,
-                    *intermediate_size,
-                    hidden_size,
-                    device=tensor.device,
-                    dtype=tensor.dtype,
-                    requires_grad=tensor.requires_grad,
-                )
+                gathered_tensor = MemoryBuffer().get("allgather", (gathered_batch_size, *intermediate_size, hidden_size), dtype=tensor.dtype)
 
                 handle = dist.all_gather_into_tensor(gathered_tensor, tensor, group=group, async_op=True)
 
@@ -261,7 +254,7 @@ def backward(ctx, grad_output):
         use_bias = ctx.use_bias
         tp_mode = ctx.tp_mode
 
-        handle: Optional[dist.Work] = None
+        handle1: Optional[dist.Work] = None
         if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
             # TODO @thomasw21: gather along another dimension
             sharded_batch_size, *rest_size = tensor.shape
@@ -273,14 +266,8 @@ def backward(ctx, grad_output):
             else:
                 unsharded_batch_size = sharded_batch_size * group.size()
 
-                unsharded_tensor = torch.empty(
-                    unsharded_batch_size,
-                    *rest_size,
-                    device=tensor.device,
-                    dtype=tensor.dtype,
-                    requires_grad=False,
-                )
-                handle = dist.all_gather_into_tensor(unsharded_tensor, tensor, group=group, async_op=True)
+                unsharded_tensor = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=tensor.dtype)
+                handle1 = dist.all_gather_into_tensor(unsharded_tensor, tensor, group=group, async_op=True)
                 # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
                 # gather is scheduled before the tensor gradient computation
                 total_tensor = unsharded_tensor
@@ -289,9 +276,6 @@ def backward(ctx, grad_output):
 
         grad_tensor = grad_output.matmul(weight)
 
-        if handle is not None:
-            handle.wait()
-
         # Doing gather + slicing during the NeMo forward pass can make this tensor
         # not be contiguous. PyTorch only checks if the tensor is contiguous, and only
         # clones it if it's not contiguous:
@@ -303,7 +287,7 @@ def backward(ctx, grad_output):
         grad_output = grad_output.view(math.prod(grad_output_first_dims), grad_output_last_dim)
         total_tensor = total_tensor.view(math.prod(total_tensor_first_dims), total_tensor_last_dim)
 
-        handle: Optional[dist.Work] = None
+        handle2: Optional[dist.Work] = None
         if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
             if group.size() == 1:
                 sub_grad_tensor = grad_tensor
@@ -312,23 +296,27 @@ def backward(ctx, grad_output):
                     tensor.shape, dtype=grad_tensor.dtype, device=grad_tensor.device, requires_grad=False
                 )
                 # reduce_scatter
-                handle = dist.reduce_scatter_tensor(sub_grad_tensor, grad_tensor, group=group, async_op=True)
+                handle2 = dist.reduce_scatter_tensor(sub_grad_tensor, grad_tensor, group=group, async_op=True)
                 # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
                 # reduce scatter is scheduled before the weight gradient computation
         elif tp_mode is TensorParallelLinearMode.ALL_REDUCE:
             # Asynchronous all-reduce
-            handle = dist.all_reduce(grad_tensor, group=group, async_op=True)
+            handle2 = dist.all_reduce(grad_tensor, group=group, async_op=True)
             # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
             # all-reduce is scheduled before the weight gradient computation
         else:
             raise ValueError()
 
+        grad_bias = grad_output.sum(dim=0) if use_bias else None
+
+        if handle1 is not None:
+            handle1.wait()
+
         # TODO @thomasw21: This sounds like we don't have the optimal physical layout
         grad_weight = grad_output.t().matmul(total_tensor)
-        grad_bias = grad_output.sum(dim=0) if use_bias else None
 
-        if handle is not None:
-            handle.wait()
+        if handle2 is not None:
+            handle2.wait()
 
         if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
             return sub_grad_tensor, grad_weight, grad_bias, None, None
@@ -472,13 +460,7 @@ def backward(ctx, grad_output):
         else:
             unsharded_batch_size = sharded_batch_size * group.size()
 
-            total_grad_output = torch.empty(
-                unsharded_batch_size,
-                *rest_size,
-                device=grad_output.device,
-                dtype=grad_output.dtype,
-                requires_grad=False,
-            )
+            total_grad_output = MemoryBuffer().get("allgather2", (unsharded_batch_size, *rest_size), dtype=tensor.dtype)
 
             # Doing gather + slicing during the NeMo forward pass can make this tensor
             # not be contiguous. PyTorch only checks if the tensor is contiguous, and only

From cb0f2609e357747a0a3001dd9b12c649f9e6eef7 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Thu, 18 Jul 2024 11:17:09 +0000
Subject: [PATCH 23/33] precommit

---
 .../parallel/tensor_parallel/functional.py       | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 1a82254e..3821d544 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -19,14 +19,13 @@
 from torch.nn import functional as F
 
 import nanotron.distributed as dist
-from nanotron.parallel.utils import MemoryBuffer
 from nanotron.parallel.tensor_parallel.distributed_differentiable_primitives import (
     differentiable_all_reduce_sum,
     differentiable_identity,
     differentiable_reduce_scatter_sum,
 )
 from nanotron.parallel.tensor_parallel.enum import TensorParallelLinearMode
-from nanotron.parallel.utils import assert_cuda_max_connections_set_to_1
+from nanotron.parallel.utils import MemoryBuffer, assert_cuda_max_connections_set_to_1
 
 
 class _ShardedCrossEntropy(torch.autograd.Function):
@@ -149,7 +148,9 @@ def forward(ctx, tensor, weight, bias, group, tp_mode):
                     group = dist.distributed_c10d._get_default_group()
                 gathered_batch_size = sharded_batch_size * group.size()
 
-                gathered_tensor = MemoryBuffer().get("allgather", (gathered_batch_size, *intermediate_size, hidden_size), dtype=tensor.dtype)
+                gathered_tensor = MemoryBuffer().get(
+                    "allgather", (gathered_batch_size, *intermediate_size, hidden_size), dtype=tensor.dtype
+                )
 
                 handle = dist.all_gather_into_tensor(gathered_tensor, tensor, group=group, async_op=True)
 
@@ -266,7 +267,9 @@ def backward(ctx, grad_output):
             else:
                 unsharded_batch_size = sharded_batch_size * group.size()
 
-                unsharded_tensor = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=tensor.dtype)
+                unsharded_tensor = MemoryBuffer().get(
+                    "allgather", (unsharded_batch_size, *rest_size), dtype=tensor.dtype
+                )
                 handle1 = dist.all_gather_into_tensor(unsharded_tensor, tensor, group=group, async_op=True)
                 # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
                 # gather is scheduled before the tensor gradient computation
@@ -399,7 +402,6 @@ def backward(ctx, grad_output: torch.Tensor):
         return sub_grad_input, grad_weight, grad_bias, None, None
 
 
-
 def column_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -460,7 +462,9 @@ def backward(ctx, grad_output):
         else:
             unsharded_batch_size = sharded_batch_size * group.size()
 
-            total_grad_output = MemoryBuffer().get("allgather2", (unsharded_batch_size, *rest_size), dtype=tensor.dtype)
+            total_grad_output = MemoryBuffer().get(
+                "allgather2", (unsharded_batch_size, *rest_size), dtype=tensor.dtype
+            )
 
             # Doing gather + slicing during the NeMo forward pass can make this tensor
             # not be contiguous. PyTorch only checks if the tensor is contiguous, and only

From 6d85d038d52ffc06ec4f2ae4705deae3b05d25d8 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Thu, 18 Jul 2024 11:46:39 +0000
Subject: [PATCH 24/33] Added no_recompute_allgather mode to async

---
 .../parallel/tensor_parallel/functional.py    | 36 +++++++++++++------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 3821d544..29480f9a 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -120,10 +120,12 @@ class _ColumnLinearAsyncCommunication(torch.autograd.Function):
 
     @staticmethod
     @assert_cuda_max_connections_set_to_1
-    def forward(ctx, tensor, weight, bias, group, tp_mode):
+    def forward(ctx, tensor, weight, bias, group, tp_mode, tp_recompute_allgather):
         ctx.use_bias = bias is not None
         ctx.tp_mode = tp_mode
         ctx.group = group
+        ctx.tp_recompute_allgather = tp_recompute_allgather
+        ctx.tensor_shape = tensor.size()
 
         if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
             gathered_tensor = tensor
@@ -140,7 +142,7 @@ def forward(ctx, tensor, weight, bias, group, tp_mode):
                 # `tensor` can sometimes not be contiguous
                 # https://cs.github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L317
                 tensor = tensor.contiguous()
-                ctx.save_for_backward(tensor, weight)
+                # ctx.save_for_backward(tensor, weight)
 
                 # TODO @thomasw21: gather along another dimension
                 sharded_batch_size, *intermediate_size, hidden_size = tensor.shape
@@ -148,9 +150,19 @@ def forward(ctx, tensor, weight, bias, group, tp_mode):
                     group = dist.distributed_c10d._get_default_group()
                 gathered_batch_size = sharded_batch_size * group.size()
 
-                gathered_tensor = MemoryBuffer().get(
-                    "allgather", (gathered_batch_size, *intermediate_size, hidden_size), dtype=tensor.dtype
-                )
+                if tp_recompute_allgather:
+                    gathered_tensor = MemoryBuffer().get(
+                        "allgather", (gathered_batch_size, *intermediate_size, hidden_size), dtype=tensor.dtype
+                    )
+                else:
+                    gathered_tensor = torch.empty(
+                        gathered_batch_size,
+                        *intermediate_size,
+                        hidden_size,
+                        device=tensor.device,
+                        dtype=tensor.dtype,
+                        requires_grad=False,
+                    )
 
                 handle = dist.all_gather_into_tensor(gathered_tensor, tensor, group=group, async_op=True)
 
@@ -198,6 +210,10 @@ def forward(ctx, tensor, weight, bias, group, tp_mode):
 
                 # Wait communication
                 handle.wait()
+                if tp_recompute_allgather:
+                    ctx.save_for_backward(tensor, weight)
+                else:
+                    ctx.save_for_backward(gathered_tensor, weight)
 
                 # Compute all the other shards that are obtained from AllGather
                 # weights: w0 w1 w2 w3
@@ -256,7 +272,7 @@ def backward(ctx, grad_output):
         tp_mode = ctx.tp_mode
 
         handle1: Optional[dist.Work] = None
-        if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
+        if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER and ctx.tp_recompute_allgather:
             # TODO @thomasw21: gather along another dimension
             sharded_batch_size, *rest_size = tensor.shape
             if group is None:
@@ -296,7 +312,7 @@ def backward(ctx, grad_output):
                 sub_grad_tensor = grad_tensor
             else:
                 sub_grad_tensor = torch.empty(
-                    tensor.shape, dtype=grad_tensor.dtype, device=grad_tensor.device, requires_grad=False
+                    ctx.tensor_shape, dtype=grad_tensor.dtype, device=grad_tensor.device, requires_grad=False
                 )
                 # reduce_scatter
                 handle2 = dist.reduce_scatter_tensor(sub_grad_tensor, grad_tensor, group=group, async_op=True)
@@ -322,9 +338,9 @@ def backward(ctx, grad_output):
             handle2.wait()
 
         if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
-            return sub_grad_tensor, grad_weight, grad_bias, None, None
+            return sub_grad_tensor, grad_weight, grad_bias, None, None, None
         elif tp_mode is TensorParallelLinearMode.ALL_REDUCE:
-            return grad_tensor, grad_weight, grad_bias, None, None
+            return grad_tensor, grad_weight, grad_bias, None, None, None
         else:
             raise ValueError(f"Got unexpected mode: {tp_mode}.")
 
@@ -412,7 +428,7 @@ def column_linear(
     tp_recompute_allgather: bool = True,
 ):
     if async_communication:
-        return _ColumnLinearAsyncCommunication.apply(input, weight, bias, group, tp_mode)
+        return _ColumnLinearAsyncCommunication.apply(input, weight, bias, group, tp_mode, tp_recompute_allgather)
 
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         input = differentiable_identity(input, group=group)

From 2afd00769c7c5891341e2d7880492bfc80c524f6 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 23 Jul 2024 09:52:12 +0000
Subject: [PATCH 25/33] Fixed List not found

---
 src/nanotron/models/llama.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py
index 7ee44390..49ea86e6 100644
--- a/src/nanotron/models/llama.py
+++ b/src/nanotron/models/llama.py
@@ -14,7 +14,7 @@
 # limitations under the License.
 """PyTorch LLaMa model."""
 
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Union, List
 
 import torch
 from torch import nn

From 7e758db3068948178edd2151232e4abd7b2d5ffd Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 23 Jul 2024 13:17:34 +0000
Subject: [PATCH 26/33] Fixed tp=1 case

---
 .../parallel/tensor_parallel/functional.py    | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 2b93fb02..22c8ca3c 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -358,11 +358,14 @@ def forward(
         # Do allgather.
         sharded_batch_size, *rest_size = input.shape
         unsharded_batch_size = sharded_batch_size * group.size()
-        if tp_recompute_allgather:
+        if group.size() == 1:
+            total_input = input.contiguous()
+        elif tp_recompute_allgather:
             total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
+            dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
         else:
             total_input = torch.empty(unsharded_batch_size, *rest_size, dtype=input.dtype, device=input.device)
-        dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
+            dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
 
         # Prepare context.
         ctx.group = group
@@ -383,21 +386,22 @@ def backward(ctx, grad_output: torch.Tensor):
         group = ctx.group
         tp_recompute_allgather = ctx.tp_recompute_allgather
         input_size = ctx.input_size
-        if tp_recompute_allgather:
+        if group.size() == 1 or not tp_recompute_allgather:
+            total_input, weight, bias = ctx.saved_tensors
+        else:
             input, weight, bias = ctx.saved_tensors
             sharded_batch_size, *rest_size = input.shape
             total_input = sharded_batch_size * group.size()
             unsharded_batch_size = sharded_batch_size * group.size()
             total_input = MemoryBuffer().get("allgather", (unsharded_batch_size, *rest_size), dtype=input.dtype)
             dist.all_gather_into_tensor(total_input, input.contiguous(), group=group)
-        else:
-            total_input, weight, bias = ctx.saved_tensors
 
-        # Get the grad_output and total_input on the correct views to be able to transpose them below.
+        # Convert the tensor shapes to 2D for execution compatibility
         grad_output = grad_output.contiguous()
-        assert grad_output.dim() == 3
-        grad_output = grad_output.view(grad_output.size(0) * grad_output.size(1), grad_output.size(2))
-        total_input = total_input.view(total_input.size(0) * total_input.size(1), total_input.size(2))
+        grad_output_first_dims, grad_output_last_dim = grad_output.shape[:-1], grad_output.shape[-1]
+        total_tensor_first_dims, total_tensor_last_dim = total_tensor.shape[:-1], total_tensor.shape[-1]
+        grad_output = grad_output.view(math.prod(grad_output_first_dims), grad_output_last_dim)
+        total_tensor = total_tensor.view(math.prod(total_tensor_first_dims), total_tensor_last_dim)
 
         # Compute gradients.
         grad_weight = grad_output.T @ total_input

From 41f11f01798b88538bbc31a793435a6572b088ea Mon Sep 17 00:00:00 2001
From: Tiancheng Chen <tiachen@student.ethz.ch>
Date: Thu, 16 May 2024 18:54:55 +0200
Subject: [PATCH 27/33] fix row parallel

---
 .../parallel/tensor_parallel/functional.py    | 74 ++++++++++++++-----
 tests/test_tensor_parallel.py                 | 29 ++++++--
 2 files changed, 76 insertions(+), 27 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index fdef48ac..1c4db5de 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -387,8 +387,7 @@ def backward(ctx, grad_output):
         group = ctx.group
         use_bias = ctx.use_bias
 
-        handle_0: Optional[dist.Work] = None
-        handle_1: Optional[dist.Work] = None
+        handle: Optional[dist.Work] = None
 
         # TODO @thomasw21: gather along another dimension
         sharded_batch_size, *rest_size = grad_output.shape
@@ -412,31 +411,69 @@ def backward(ctx, grad_output):
             # https://github.com/pytorch/pytorch/blob/c47cf9bc7f9e02f649ab4ed53fe4d35732c92ab6/torch/_refs/__init__.py#L2761
             grad_output = grad_output.contiguous()
 
-            handle_0 = dist.all_gather_into_tensor(total_grad_output, grad_output, group=group, async_op=True)
-
-        grad_tensor = grad_output.matmul(weight)
+            handle = dist.all_gather_into_tensor(total_grad_output, grad_output, group=group, async_op=True)
 
-        # wait for the first all_gather to finish before starting the second all_gather
-        if handle_0 is not None:
-            handle_0.wait()
-
-        # TODO @thomasw21: gather along another dimension
-        sharded_batch_size, *rest_size = grad_tensor.shape
+        # total_grad_output: [b, s, h_out]
+        # weight: [h_out, h_in/n]
+        # total_grad_tensor: [b, s, h_in/n]
+        # grad_output: [b/n, s, h_out]
+        sharded_batch_size, *rest_size_grad_output = grad_output.shape
+        rest_size_grad_tensor = rest_size_grad_output[:-1] + [weight.shape[1]]
 
         if group.size() == 1:
-            total_grad_tensor = grad_tensor
+            total_grad_tensor = grad_output.matmul(weight)
         else:
             unsharded_batch_size = sharded_batch_size * group.size()
-
             total_grad_tensor = torch.empty(
                 unsharded_batch_size,
-                *rest_size,
-                device=grad_tensor.device,
-                dtype=grad_tensor.dtype,
+                *rest_size_grad_tensor,
+                device=grad_output.device,
+                dtype=grad_output.dtype,
                 requires_grad=False,
             )
+            before_shard_grad_tensor, same_device_shard_grad_tensor, after_shard_grad_tensor = torch.split(
+                total_grad_tensor,
+                split_size_or_sections=[
+                    sharded_batch_size * dist.get_rank(group),
+                    sharded_batch_size,
+                    sharded_batch_size * (group.size() - dist.get_rank(group) - 1),
+                ],
+                dim=0,
+            )
+            # compute local shard
+            torch.mm(
+                input=grad_output.view(-1, grad_output.shape[-1]),
+                mat2=weight,
+                out=same_device_shard_grad_tensor.view(-1, weight.shape[1]),
+            )
 
-            handle_1 = dist.all_gather_into_tensor(total_grad_tensor, grad_tensor, group=group, async_op=True)
+            if handle is not None:
+                handle.wait()
+
+            before_shard_grad_output, _, after_shard_grad_output = torch.split(
+                total_grad_output,
+                split_size_or_sections=[
+                    sharded_batch_size * dist.get_rank(group),
+                    sharded_batch_size,
+                    sharded_batch_size * (group.size() - dist.get_rank(group) - 1),
+                ],
+                dim=0,
+            )
+
+            # before shard compute
+            if before_shard_grad_tensor.numel() > 0:
+                torch.mm(
+                    input=before_shard_grad_output.view(-1, before_shard_grad_output.shape[-1]),
+                    mat2=weight,
+                    out=before_shard_grad_tensor.view(-1, weight.shape[1]),
+                )
+            # after shard compute
+            if after_shard_grad_tensor.numel() > 0:
+                torch.mm(
+                    input=after_shard_grad_output.view(-1, after_shard_grad_output.shape[-1]),
+                    mat2=weight,
+                    out=after_shard_grad_tensor.view(-1, weight.shape[1]),
+                )
 
         # Convert the tensor shapes to 2D for execution compatibility
         tensor = tensor.contiguous()
@@ -454,9 +491,6 @@ def backward(ctx, grad_output):
         grad_weight = total_grad_output.t().matmul(tensor)
         grad_bias = total_grad_output.sum(dim=0) if use_bias else None
 
-        if handle_1 is not None:
-            handle_1.wait()
-
         return total_grad_tensor, grad_weight, grad_bias, None, None
 
 
diff --git a/tests/test_tensor_parallel.py b/tests/test_tensor_parallel.py
index 127ba2fa..f5dcaeb0 100644
--- a/tests/test_tensor_parallel.py
+++ b/tests/test_tensor_parallel.py
@@ -208,14 +208,19 @@ def _test_row_linear(parallel_context: ParallelContext, tp_mode: TensorParallelL
     random_input = torch.randn(batch_size, in_features, device="cuda")
     # synchronize random_input across tp
     dist.all_reduce(random_input, op=dist.ReduceOp.AVG, group=parallel_context.tp_pg)
-
+    random_input.requires_grad = True
     # Row linear receives as input sharded input
-    random_sharded_input = random_input[
-        :,
-        dist.get_rank(parallel_context.tp_pg)
-        * in_features_per_rank : (dist.get_rank(parallel_context.tp_pg) + 1)
-        * in_features_per_rank,
-    ]
+    random_sharded_input = (
+        random_input[
+            :,
+            dist.get_rank(parallel_context.tp_pg)
+            * in_features_per_rank : (dist.get_rank(parallel_context.tp_pg) + 1)
+            * in_features_per_rank,
+        ]
+        .detach()
+        .clone()
+    )
+    random_sharded_input.requires_grad = True
 
     # Test that we get the same output after forward pass
     # TODO @kunhao: We may want to have our custom error type
@@ -261,6 +266,16 @@ def _test_row_linear(parallel_context: ParallelContext, tp_mode: TensorParallelL
     else:
         assert row_linear.bias is None
 
+    torch.testing.assert_close(
+        random_sharded_input.grad,
+        random_input.grad[
+            :,
+            dist.get_rank(parallel_context.tp_pg)
+            * in_features_per_rank : (dist.get_rank(parallel_context.tp_pg) + 1)
+            * in_features_per_rank,
+        ],
+    )
+
     parallel_context.destroy()
 
 

From cd84d4fa7bff4017a9b1653c4b68290c00eae649 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 30 Jul 2024 19:16:14 +0200
Subject: [PATCH 28/33] Fixed column parallel

---
 .../parallel/tensor_parallel/functional.py    | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 054d41e9..468855a5 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -19,14 +19,13 @@
 from torch.nn import functional as F
 
 import nanotron.distributed as dist
-from nanotron.parallel.utils import MemoryBuffer
 from nanotron.parallel.tensor_parallel.distributed_differentiable_primitives import (
     differentiable_all_reduce_sum,
     differentiable_identity,
     differentiable_reduce_scatter_sum,
 )
 from nanotron.parallel.tensor_parallel.enum import TensorParallelLinearMode
-from nanotron.parallel.utils import assert_cuda_max_connections_set_to_1
+from nanotron.parallel.utils import MemoryBuffer, assert_cuda_max_connections_set_to_1
 
 
 class _ShardedCrossEntropy(torch.autograd.Function):
@@ -399,23 +398,25 @@ def backward(ctx, grad_output: torch.Tensor):
         # Convert the tensor shapes to 2D for execution compatibility
         grad_output = grad_output.contiguous()
         grad_output_first_dims, grad_output_last_dim = grad_output.shape[:-1], grad_output.shape[-1]
-        total_tensor_first_dims, total_tensor_last_dim = total_tensor.shape[:-1], total_tensor.shape[-1]
+        total_input_first_dims, total_input_last_dim = total_input.shape[:-1], total_input.shape[-1]
         grad_output = grad_output.view(math.prod(grad_output_first_dims), grad_output_last_dim)
-        total_tensor = total_tensor.view(math.prod(total_tensor_first_dims), total_tensor_last_dim)
+        total_input = total_input.view(math.prod(total_input_first_dims), total_input_last_dim)
 
         # Compute gradients.
         grad_weight = grad_output.T @ total_input
         grad_input = grad_output @ weight
-        sub_grad_input = torch.empty(
-            input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False
-        )
-        dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
+        if group.size() == 1:
+            sub_grad_input = grad_input
+        else:
+            sub_grad_input = torch.empty(
+                input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False
+            )
+            dist.reduce_scatter_tensor(sub_grad_input, grad_input, group=group, op=dist.ReduceOp.SUM)
         grad_bias = torch.sum(grad_output, dim=0) if bias is not None else None
 
         return sub_grad_input, grad_weight, grad_bias, None, None
 
 
-
 def column_linear(
     input: torch.Tensor,
     weight: torch.Tensor,

From d3db06acb2e3fe235ae512861ff64ee9fbc9ac11 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 30 Jul 2024 19:16:28 +0200
Subject: [PATCH 29/33] Added tp_recompute_allgather test

---
 tests/test_tensor_parallel.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tests/test_tensor_parallel.py b/tests/test_tensor_parallel.py
index f5dcaeb0..8e73973b 100644
--- a/tests/test_tensor_parallel.py
+++ b/tests/test_tensor_parallel.py
@@ -18,17 +18,30 @@
 @pytest.mark.parametrize("tp,dp,pp", [pytest.param(i, 1, 1) for i in range(1, min(4, available_gpus()) + 1)])
 @pytest.mark.parametrize("tp_mode", list(TensorParallelLinearMode))
 @pytest.mark.parametrize("async_communication", [False, True])
+@pytest.mark.parametrize("tp_recompute_allgather", [False, True])
 @rerun_if_address_is_in_use()
-def test_column_linear(tp: int, dp: int, pp: int, tp_mode: TensorParallelLinearMode, async_communication: bool):
+def test_column_linear(
+    tp: int,
+    dp: int,
+    pp: int,
+    tp_mode: TensorParallelLinearMode,
+    async_communication: bool,
+    tp_recompute_allgather: bool,
+):
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE and async_communication:
         pytest.skip("ALL_REDUCE mode does not support async communication")
+    if tp_mode is TensorParallelLinearMode.ALL_REDUCE and tp_recompute_allgather:
+        pytest.skip("ALL_REDUCE mode is unaffected by tp_recompute_allgather")
     init_distributed(tp=tp, dp=dp, pp=pp)(_test_column_linear)(
-        tp_mode=tp_mode, async_communication=async_communication
+        tp_mode=tp_mode, async_communication=async_communication, tp_recompute_allgather=tp_recompute_allgather
     )
 
 
 def _test_column_linear(
-    parallel_context: ParallelContext, tp_mode: TensorParallelLinearMode, async_communication: bool
+    parallel_context: ParallelContext,
+    tp_mode: TensorParallelLinearMode,
+    async_communication: bool,
+    tp_recompute_allgather: bool,
 ):
     if async_communication:
         os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
@@ -44,6 +57,7 @@ def _test_column_linear(
         mode=tp_mode,
         device="cuda",
         async_communication=async_communication,
+        tp_recompute_allgather=tp_recompute_allgather,
     )
 
     # Un-sharded
@@ -86,7 +100,7 @@ def _test_column_linear(
             random_input = sharded_random_input
     else:
         ValueError(f"Unsupported mode: {tp_mode}")
-    # It's important that `random_input` and `sharded_random_input` are two seperate tensors with seperate storage
+    # It's important that `random_input` and `sharded_random_input` are two separate tensors with separate storage
     sharded_random_input = sharded_random_input.clone()
     random_input.requires_grad = True
     sharded_random_input.requires_grad = True

From 4c94b99a8cafd5f8dc2b7f208341a17a4d818234 Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Tue, 30 Jul 2024 19:34:31 +0200
Subject: [PATCH 30/33] Added tp_recompute_allgather test

---
 tests/test_tensor_parallel.py | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/tests/test_tensor_parallel.py b/tests/test_tensor_parallel.py
index 8e73973b..16008eaa 100644
--- a/tests/test_tensor_parallel.py
+++ b/tests/test_tensor_parallel.py
@@ -164,15 +164,32 @@ def _test_column_linear(
 @pytest.mark.parametrize("tp,dp,pp", [pytest.param(i, 1, 1) for i in range(1, min(4, available_gpus()) + 1)])
 @pytest.mark.parametrize("tp_mode", list(TensorParallelLinearMode))
 @pytest.mark.parametrize("async_communication", [False, True])
+@pytest.mark.parametrize("tp_recompute_allgather", [False, True])
 @rerun_if_address_is_in_use()
-def test_row_linear(tp: int, dp: int, pp: int, tp_mode: TensorParallelLinearMode, async_communication: bool):
+def test_row_linear(
+    tp: int,
+    dp: int,
+    pp: int,
+    tp_mode: TensorParallelLinearMode,
+    async_communication: bool,
+    tp_recompute_allgather: bool,
+):
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE and async_communication:
         pytest.skip("ALL_REDUCE mode does not support async communication")
+    if tp_mode is TensorParallelLinearMode.ALL_REDUCE and tp_recompute_allgather:
+        pytest.skip("ALL_REDUCE mode is not affected by tp_recompute_allgather")
 
-    init_distributed(tp=tp, dp=dp, pp=pp)(_test_row_linear)(tp_mode=tp_mode, async_communication=async_communication)
+    init_distributed(tp=tp, dp=dp, pp=pp)(_test_row_linear)(
+        tp_mode=tp_mode, async_communication=async_communication, tp_recompute_allgather=tp_recompute_allgather
+    )
 
 
-def _test_row_linear(parallel_context: ParallelContext, tp_mode: TensorParallelLinearMode, async_communication: bool):
+def _test_row_linear(
+    parallel_context: ParallelContext,
+    tp_mode: TensorParallelLinearMode,
+    async_communication: bool,
+    tp_recompute_allgather: bool,
+):
     if async_communication:
         os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
     out_features = 3

From 3c3561158eb053176b6f148a2366ea1aa56fdc7d Mon Sep 17 00:00:00 2001
From: Phuc Nguyen <b3f0cus@icloud.com>
Date: Wed, 31 Jul 2024 17:15:34 +0000
Subject: [PATCH 31/33] change to correct config_nanoset.yaml path

---
 docs/nanoset.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/nanoset.md b/docs/nanoset.md
index 9dce21b7..61393438 100644
--- a/docs/nanoset.md
+++ b/docs/nanoset.md
@@ -79,7 +79,7 @@ To work with `Nanosets`, we just need to configure 1 argument:
 
 Finally, to use the `Nanosets`, launch the training with [`run_train.py`](../run_train.py).
 ```shell
-torchrun --nproc-per-node 8 run_train.py --config configs/config_nanoset.yaml
+torchrun --nproc-per-node 1 run_train.py --config examples/config_nanoset.yaml
 ```
 
 ## Under the hood

From 7daa186e84e03fa66fa83129d9b0acffb1a668ba Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Fri, 2 Aug 2024 15:40:46 +0200
Subject: [PATCH 32/33] Minor restyling

---
 .../parallel/tensor_parallel/functional.py        | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 468855a5..1fb86cb5 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -115,7 +115,7 @@ def sharded_cross_entropy(sharded_logits, target, group: dist.ProcessGroup, dtyp
     return _ShardedCrossEntropy.apply(sharded_logits, target, group)
 
 
-class _ColumnLinearAsyncCommunication(torch.autograd.Function):
+class _ColumnLinearNoAsyncCommunicationReduceScatterMode(torch.autograd.Function):
     """Adapted from https://github.com/NVIDIA/Megatron-LM/blob/e6d7e09845590d0a36bc7f29eb28db974fb8da4e/megatron/core/tensor_parallel/layers.py#L215"""
 
     @staticmethod
@@ -408,6 +408,9 @@ def backward(ctx, grad_output: torch.Tensor):
         if group.size() == 1:
             sub_grad_input = grad_input
         else:
+            # Seems that `reduce_scatter` need contiguous tensors: https://github.com/pytorch/pytorch/blob/2b267fa7f28e18ca6ea1de4201d2541a40411457/torch/distributed/nn/functional.py#L305
+            # We set grad_input to be contiguous in case it isn't already.
+            grad_input = grad_input.contiguous()
             sub_grad_input = torch.empty(
                 input_size, dtype=total_input.dtype, device=total_input.device, requires_grad=False
             )
@@ -427,16 +430,14 @@ def column_linear(
     tp_recompute_allgather: bool = True,
 ):
     if async_communication:
-        return _ColumnLinearAsyncCommunication.apply(input, weight, bias, group, tp_mode)
+        return _ColumnLinearNoAsyncCommunicationReduceScatterMode.apply(input, weight, bias, group, tp_mode)
 
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         input = differentiable_identity(input, group=group)
-    elif tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
+        return F.linear(input, weight, bias)
+    if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
         return _ColumnLinearContextParallelNoAsync.apply(input, weight, bias, group, tp_recompute_allgather)
-    else:
-        raise ValueError(f"Got unexpected mode: {tp_mode}.")
-
-    return F.linear(input, weight, bias)
+    raise ValueError(f"Got unexpected mode: {tp_mode}.")
 
 
 class _RowLinearAsyncCommunication(torch.autograd.Function):

From 31c3c5ad0a845ff6318842c663223e9621586a3d Mon Sep 17 00:00:00 2001
From: AleHD <alejandro.hernandezcano@epfl.ch>
Date: Fri, 2 Aug 2024 15:54:44 +0200
Subject: [PATCH 33/33] Fixed names

---
 src/nanotron/parallel/tensor_parallel/functional.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/nanotron/parallel/tensor_parallel/functional.py b/src/nanotron/parallel/tensor_parallel/functional.py
index 1fb86cb5..7a88aec6 100644
--- a/src/nanotron/parallel/tensor_parallel/functional.py
+++ b/src/nanotron/parallel/tensor_parallel/functional.py
@@ -115,7 +115,7 @@ def sharded_cross_entropy(sharded_logits, target, group: dist.ProcessGroup, dtyp
     return _ShardedCrossEntropy.apply(sharded_logits, target, group)
 
 
-class _ColumnLinearNoAsyncCommunicationReduceScatterMode(torch.autograd.Function):
+class _ColumnLinearAsyncCommunication(torch.autograd.Function):
     """Adapted from https://github.com/NVIDIA/Megatron-LM/blob/e6d7e09845590d0a36bc7f29eb28db974fb8da4e/megatron/core/tensor_parallel/layers.py#L215"""
 
     @staticmethod
@@ -337,7 +337,7 @@ def backward(ctx, grad_output):
             raise ValueError(f"Got unexpected mode: {tp_mode}.")
 
 
-class _ColumnLinearContextParallelNoAsync(torch.autograd.Function):
+class _ColumnLinearNoAsyncCommunicationReduceScatterMode(torch.autograd.Function):
     """
     Column linear with memory_buffer for the allgather, context parallel
     enabled (i.e. tp_mode = TensorParallelLinearMode.REDUCE_SCATTER) and
@@ -430,13 +430,15 @@ def column_linear(
     tp_recompute_allgather: bool = True,
 ):
     if async_communication:
-        return _ColumnLinearNoAsyncCommunicationReduceScatterMode.apply(input, weight, bias, group, tp_mode)
+        return _ColumnLinearAsyncCommunication.apply(input, weight, bias, group, tp_mode)
 
     if tp_mode is TensorParallelLinearMode.ALL_REDUCE:
         input = differentiable_identity(input, group=group)
         return F.linear(input, weight, bias)
     if tp_mode is TensorParallelLinearMode.REDUCE_SCATTER:
-        return _ColumnLinearContextParallelNoAsync.apply(input, weight, bias, group, tp_recompute_allgather)
+        return _ColumnLinearNoAsyncCommunicationReduceScatterMode.apply(
+            input, weight, bias, group, tp_recompute_allgather
+        )
     raise ValueError(f"Got unexpected mode: {tp_mode}.")