Distributed dataloader

chemlactica/config/create_train_config.py

@@ -20,6 +20,9 @@
 model_train_configs["125m"][
     "tokenizer_path"
 ] = "chemlactica/tokenizer/ChemLacticaTokenizer66"
+model_train_configs["small_opt"][
+    "tokenizer_path"
+] = "chemlactica/tokenizer/ChemLacticaTokenizer66"
 model_train_configs["1.3b"][
     "tokenizer_path"
 ] = "chemlactica/tokenizer/ChemLacticaTokenizer66"

chemlactica/jsonl_dataset.py

@@ -1,8 +1,9 @@
 from typing import List
-import torch
 
-# from io import StringIO
 import os
+from accelerate.state import PartialState
+
+distributed_state = PartialState()
 
 
 def generator_init_print(shared_jsonl_files, files):
@@ -23,51 +24,39 @@ def setup_generator(shared_jsonl_files, files):
     return file_states
 
 
-def get_batch(file, state, chunk_size):
-    with open(file) as f:
-        f.seek(state["position"])
-        batch = f.read(chunk_size)
-        if not batch:
-            raise StopIteration
-
-        batch += f.readline()
-        batch = batch.splitlines()
+def should_yield_on_current_rank(i, num_processes, process_index):
+    return i % num_processes == process_index
 
-        # batch = [line.rstrip("\n") for line in batch]
-        state["position"] = f.tell()
-        batch_len = len(batch)
-        state["line_number"] += batch_len
-    return batch, batch_len, state
 
-
-def format_sample(sample, return_line_info, batch_len, file, state, i):
+def format_sample(line):
+    sample = line.strip()
     ret = {"text": sample}
-    if return_line_info:
-        ret["line_info"] = {
-            "file": file,
-            "line_number": state["line_number"] - batch_len + i,
-        }
     return ret
 
 
 def samples_generator(
     files: List[str], shared_jsonl_files, chunk_size=25000, return_line_info=False
 ):
-    if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
-        file_states = setup_generator(shared_jsonl_files, files)
-
-        returned = True
-        while returned:
-            returned = False
-            for file, state in file_states.items():
-                try:
-                    batch, batch_len, state = get_batch(file, state, chunk_size)
-                except StopIteration:
-                    break
-                for i, sample in enumerate(batch, start=1):
-                    returned = True
-                    ret = format_sample(
-                        sample, return_line_info, batch_len, file, state, i
-                    )
-                    yield ret
-                shared_jsonl_files[file] = state
+    file_states = setup_generator(shared_jsonl_files, files)
+
+    returned = True
+    while returned:
+        returned = False
+        for file, state in file_states.items():
+            with open(file) as f:
+                f.seek(state["position"])
+                line = f.readline()
+                counter = 0
+                while line:
+                    state["position"] = f.tell()
+                    if should_yield_on_current_rank(
+                        counter,
+                        distributed_state.num_processes,
+                        distributed_state.process_index,
+                    ):
+                        returned = True
+                        ret = format_sample(line)
+                        yield ret
+                    counter = counter + 1
+                    shared_jsonl_files[file] = state
+                    line = f.readline()

chemlactica/train.py

@@ -50,7 +50,6 @@
 
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "caching_allocator"
 # os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
 # signal.signal(signal.SIGINT, signal_handler)
 # signal.signal(signal.SIGTERM, signal_handler)
 
@@ -192,9 +191,9 @@ def train(
 
     accelerator.wait_for_everyone()
 
-    with multiprocessing.Manager() if accelerator.is_main_process else nullcontext() as manager:
+    with multiprocessing.Manager() as manager:
         shared_jsonl_files = None
-        if accelerator.is_main_process and train_type == "pretrain":
+        if train_type == "pretrain":
             shared_jsonl_files = manager.dict()
             trainer_callback_dict[
                 "json_dataset_resume_callback"

chemlactica/utils/model_utils.py

@@ -105,7 +105,7 @@ def load_model(
                 ffn_dim=model_config["ffn_dim"],
                 max_position_embeddings=model_config["max_position_embeddings"],
                 num_attention_heads=model_config["num_attention_heads"],
-                word_embed_proj_dim=model_config["word_sembed_proj_dim"],
+                word_embed_proj_dim=model_config["word_embed_proj_dim"],
             )
         )
     if "galactica" in from_pretrained.lower():