fix: config

.gitignore

@@ -1,3 +1,6 @@
+*.pkl
+ckpts*
+.deepspeed_env
 *.jsonl
 *tar.gz
 ckpts**

clean.py

@@ -64,6 +64,7 @@ for file in glob.glob(os.path.join(prompt_generation_dir, "*.jsonl")):
     df = df.dropna(subset=['prompt', 'response'])
     df = df[df['prompt'] != '']
     df = df[df['response'] != '']
+    df = df[df["prompt"].str.len() > 1]
     curr_len = len(df)
 
     print(f"Removed {prev_len - curr_len} rows")

configs/deepspeed/ds_config_gptj.json

@@ -0,0 +1,28 @@
+{
+	"train_batch_size": "auto",
+	"gradient_accumulation_steps": "auto",
+	"train_micro_batch_size_per_gpu": "auto",
+	"fp16": {
+	  "enabled": "auto",
+	  "min_loss_scale": 1,
+	  "loss_scale_window": 1000,
+	  "hysteresis": 2,
+	  "initial_scale_power": 32
+	},
+	"bf16": {
+		"enabled": "auto"
+	},
+	"gradient_clipping": 1.0,
+	"zero_optimization": {
+	  "stage": 2,
+	  "offload_param": {
+		"device": "none"
+	  },
+	  "offload_optimizer": {
+		"device": "none"
+	  },
+	  "allgather_partitions": true,
+	  "allgather_bucket_size": 5e8,
+	  "contiguous_gradients": true
+	}
+  }

configs/deepspeed/ds_config_gptj_lora.json

@@ -0,0 +1,48 @@
+{
+	"train_batch_size": "auto",
+	"gradient_accumulation_steps": "auto",
+	"train_micro_batch_size_per_gpu": "auto",
+	"fp16": {
+	  "enabled": "auto",
+	  "min_loss_scale": 1,
+	  "loss_scale_window": 1000,
+	  "hysteresis": 2,
+	  "initial_scale_power": 32
+	},
+	"bf16": {
+		"enabled": "auto"
+	},
+	"gradient_clipping": 1,
+	"zero_optimization": {
+	  "stage": 2,
+	  "offload_param": {
+		"device": "cpu"
+	  },
+	  "offload_optimizer": {
+		"device": "cpu"
+	  },
+	  "allgather_partitions": true,
+	  "allgather_bucket_size": 5e8,
+	  "contiguous_gradients": true
+	},
+	"optimizer": {
+	  "type": "AdamW",
+	  "params": {
+		"lr": "auto",
+		"betas": [
+		  0.9,
+		  0.999
+		],
+		"eps": 1e-08
+	  }
+	},
+	"scheduler": {
+	  "type": "WarmupLR",
+	  "params": {
+		"warmup_min_lr": 0,
+		"warmup_max_lr": "auto",
+		"warmup_num_steps": "auto",
+		"warmup_type": "linear"
+	  }
+	}
+  }

configs/inference/gptj.yaml

@@ -0,0 +1,14 @@
+# model/tokenizer
+model_name: "nomic-ai/gpt4all-gptj-multinode-deepspeed-finetuned-epoch_0"
+tokenizer_name: "EleutherAI/gpt-j-6B"
+
+# dataset
+streaming: false
+num_proc: 64
+dataset_path: "data_multiplus" 
+max_length: 1024
+batch_size: 32 
+
+# logging
+seed: 42
+

configs/train/finetune_gptj.yaml

@@ -0,0 +1,33 @@
+# model/tokenizer
+model_name: "EleutherAI/gpt-j-6B"
+tokenizer_name: "EleutherAI/gpt-j-6B"
+gradient_checkpointing: true
+save_name: "nomic-ai/gpt4all-mosaic"
+
+# dataset
+streaming: false
+num_proc: 64
+dataset_path: "nomic-ai/turbo-500k-multi"
+max_length: 1024
+batch_size: 8
+
+# train dynamics
+lr: 2.0e-5
+min_lr: 0 
+weight_decay: 0.0
+eval_every: 500
+eval_steps: 105
+save_every: 500
+log_grads_every: 500
+output_dir: "ckpts/gpt4all-gptj-multinode"
+checkpoint: null
+lora: false
+warmup_steps: 500
+num_epochs: 2 
+
+# logging
+wandb: true
+wandb_entity: vicuna
+wandb_project_name: vicuna
+seed: 42
+

configs/train/finetune_gptj_lora.yaml

@@ -0,0 +1,33 @@
+# model/tokenizer
+model_name: "EleutherAI/gpt-j-6b"
+tokenizer_name: "EleutherAI/gpt-j-6b"
+gradient_checkpointing: false
+save_name: "nomic-ai/gpt4all-mosaic"
+
+# dataset
+streaming: false
+num_proc: 64
+dataset_path: "nomic-ai/turbo-500k-multi"
+max_length: 1024
+batch_size: 2
+
+# train dynamics
+lr: 2.0e-5
+min_lr: 0 
+weight_decay: 0.0
+eval_every: 500
+eval_steps: 105
+save_every: 500
+log_grads_every: 500
+output_dir: "ckpts/gpt4all-gptj-multinode"
+checkpoint: null
+lora: true
+warmup_steps: 500
+num_epochs: 2 
+
+# logging
+wandb: true
+wandb_entity: zanussbaum
+wandb_project_name: mosaic
+seed: 42
+

configs/train/finetune_lora.yaml

@@ -2,17 +2,19 @@
 model_name: # update
 tokenizer_name: # update
 gradient_checkpointing: false
-save_name: "nomic-ai/gpt4all-lora-multi-turn"
+save_name: "nomic-ai/gpt4all-lora-llama"
 
 # dataset
 streaming: false
 num_proc: 64
-dataset_path: "data_multiturn"
+dataset_path: "nomic-ai/turbo-500k-multi"
 max_length: 1024
 batch_size: 4
 
 # train dynamics
 lr: 5.0e-5
+min_lr: 0
+weight_decay: 0.0
 eval_every: 2000
 eval_steps: 100
 save_every: 2000

create_hostname.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+export WORKER_IP=$1
+N_GPUS=8
+# create dir if doesn't exist
+sudo mkdir -p /job
+printf "localhost slots=$N_GPUS\n$WORKER_IP slots=$N_GPUS" | sudo tee /job/hostfile
+echo /job/hostfile

data.py

@@ -9,44 +9,49 @@ from transformers import DefaultDataCollator
 
 def tokenize_inputs(config, tokenizer, examples):
     max_length = config["max_length"]
-    input_ids = torch.full((len(examples["prompt"]), max_length), tokenizer.pad_token_id)
-    # ignore bos
-    newline_tokens = tokenizer("\n", return_tensors="pt")["input_ids"][0, 1:]
 
-    out = {"labels": [], "attention_mask": []}
-    for i, (prompt, response) in enumerate(zip(examples["prompt"], examples["response"])):
-        input_tokens = tokenizer(prompt, truncation=True, max_length=max_length // 2, return_tensors="pt")["input_ids"].squeeze()
-        input_len = len(input_tokens)
+    # hacky backward compatible
+    different_eos = tokenizer.eos_token != "</s>"
+    out = {"labels": [], "input_ids": []}
+    for prompt, response in zip(examples["prompt"], examples["response"]):
+        if different_eos:
+            if response.count("</s>") > 0:
+                response = response.replace("</s>", tokenizer.eos_token)
 
-        # plus one since we remove bos from response
-        # but we subtract one since we want to add eos token
-        remaining_tokens = max_length - input_len - len(newline_tokens) + 1
-        # remove bos
-        target_tokens = tokenizer(response, truncation=True, max_length=remaining_tokens, return_tensors="pt")["input_ids"].squeeze()[1:]
+        prompt_len = len(tokenizer(prompt + "\n", return_tensors="pt")["input_ids"][0])
 
-        input_ids[i, :input_len] = input_tokens
-        # add newline between prompt and response
-        newline_plus_inputs = input_len + len(newline_tokens)
-        input_ids[i, input_len: newline_plus_inputs] = newline_tokens
+        # hack if our prompt is super long
+        # we need to include some labels so we arbitrarily trunacate at max_length // 2
+        # if the length is too long
+        if prompt_len >= max_length // 2:
+            # if prompt is too long, truncate
+            # but make sure to truncate to at max 1024 tokens
+            new_len = min(max_length // 2, len(prompt) // 2)
+            prompt = prompt[:new_len]
+            # get new prompt length
+            prompt_len = tokenizer(prompt + "\n", return_tensors="pt", max_length=max_length // 2, truncation=True).input_ids.ne(tokenizer.pad_token_id).sum().item()
 
-        # add target tokens, remove bos
-        input_ids[i, newline_plus_inputs: newline_plus_inputs + len(target_tokens)] = target_tokens
-        # add eos token, enforce stopping if we don't truncate 
-        # we don't want long code to stop generating if truncated during training
-        if newline_plus_inputs + len(target_tokens) < max_length:
-            input_ids[i, newline_plus_inputs + len(target_tokens)] = tokenizer.eos_token_id
+        assert prompt_len <= max_length // 2, f"prompt length {prompt_len} exceeds max length {max_length}"
 
-        labels = input_ids[i].clone()
-        labels[: newline_plus_inputs] = -100
-        labels[labels == tokenizer.pad_token_id] = -100
-        # to debug this, can set all values == -100 to the pad token, then assert that tokenizer.decode(labels, skip_special_tokens=True).strip() == response
+        input_tokens = tokenizer(prompt + "\n" + response + tokenizer.eos_token,
+                                 truncation=True, max_length=max_length, return_tensors="pt")["input_ids"].squeeze()
 
-        attention_mask = input_ids[i].ne(tokenizer.pad_token_id).int()
+        labels = input_tokens.clone()
+        labels[:prompt_len] = -100
+        if len(labels) < max_length:
+            # pad to max_length with -100
+            labels = torch.cat([labels, torch.full((max_length - len(labels),), -100)])
 
+        assert (labels == -100).sum() < len(labels), f"Labels are all -100, something wrong. prompt length {prompt_len} exceeds max length {max_length}" 
+        
+        if (labels == -100).sum() == len(labels) - 1:
+            print(prompt)
+            print(response)
+            raise
+
+        input_tokens = tokenizer.pad({"input_ids": input_tokens}, padding="max_length", max_length=max_length)["input_ids"]
         out["labels"].append(labels)
-        out["attention_mask"].append(attention_mask)
-
-    out["input_ids"] = input_ids
+        out["input_ids"].append(input_tokens)
 
     out = {k: torch.stack(v) if isinstance(v, list) else v for k, v in out.items()}
 
@@ -68,7 +73,7 @@ def load_data(config, tokenizer):
         dataset = load_dataset("json", data_files=files, split="train")
 
     else:
-        dataset = load_dataset(dataset_path)
+        dataset = load_dataset(dataset_path, split="train")
 
     dataset = dataset.train_test_split(test_size=.05, seed=config["seed"])
 
@@ -111,3 +116,53 @@ def load_data(config, tokenizer):
     )
 
     return train_dataloader, val_dataloader
+
+    
+def load_data_for_inference(config, tokenizer):
+    dataset_path = config["dataset_path"]
+
+    if os.path.exists(dataset_path):
+        # check if path is a directory
+        if os.path.isdir(dataset_path):
+            files = glob.glob(os.path.join(dataset_path, "*_clean.jsonl"))
+        else:
+            files = [dataset_path]
+
+        print(f"Reading files {files}")
+
+        dataset = load_dataset("json", data_files=files, split="train")
+
+    else:
+        dataset = load_dataset(dataset_path, split="train")
+
+    dataset = dataset.train_test_split(test_size=.05, seed=config["seed"])
+
+    train_dataset, val_dataset = dataset["train"], dataset["test"]
+
+    train_dataset = train_dataset.add_column("index", list(range(len(train_dataset))))
+    # select first N batches that are divisible by batch_size
+    # gather is a bit annoying (or the way I'm using it) to get uneven batches as it duplicates data
+    train_dataset = train_dataset.select(range((len(train_dataset) // config["batch_size"]) * config["batch_size"]))
+    val_dataset = val_dataset.add_column("index", list(range(len(val_dataset))))
+    val_dataset = val_dataset.select(range((len(val_dataset) // config["batch_size"]) * config["batch_size"]))
+
+    if config["streaming"] is False:
+        kwargs = {"num_proc": config["num_proc"]}
+    else:
+        kwargs = {}
+
+    # tokenize inputs and return labels and attention mask
+    train_dataset = train_dataset.map(
+        lambda ele: tokenize_inputs(config, tokenizer, ele),
+        batched=True,
+        **kwargs
+    )
+    val_dataset = val_dataset.map(
+        lambda ele: tokenize_inputs(config, tokenizer, ele), 
+        batched=True,
+        **kwargs
+    )
+    train_dataset = train_dataset.with_format("torch")
+    val_dataset = val_dataset.with_format("torch")
+
+    return train_dataset, val_dataset

head_node_setup.sh

@@ -0,0 +1,19 @@
+#!/bin/sh
+WORKER_IP=$1
+N_GPUS=$2
+
+sudo apt install -y nfs-kernel-server
+sudo mkdir -p ./data_multiplus
+sudo chmod 777 ./data_multiplus
+printf "${PWD}/data_multiplus ${WORKER_IP}(rw,sync,no_subtree_check)" | sudo tee -a /etc/exports
+sudo systemctl restart nfs-kernel-server
+
+sudo apt-get install -y pdsh
+export DSHPATH=$PATH
+export PDSH_RCMD_TYPE=ssh
+
+ssh-keygen -t rsa -N ''
+cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys
+
+sudo mkdir -p /job
+printf "localhost slots=$N_GPUS\n$WORKER_IP slots=$N_GPUS" | sudo tee /job/hostfile

inference.py

@@ -0,0 +1,207 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import torch.nn as nn
+from argparse import ArgumentParser
+from read import read_config
+from accelerate.utils import  set_seed
+from data import load_data_for_inference
+from tqdm import tqdm
+from datasets import  Dataset
+import torch.distributed as dist
+from transformers.trainer_pt_utils import  nested_numpify
+from transformers import DefaultDataCollator
+from torch.utils.data import DataLoader, DistributedSampler
+import numpy as np
+import pyarrow as pa
+from pyarrow import compute as pc
+
+
+def calc_cross_entropy_no_reduction(lm_logits, labels):
+    # calculate cross entropy across batch dim
+    shift_logits = lm_logits[..., :-1, :].contiguous()
+    shift_labels = labels[..., 1:].contiguous()
+    # Flatten the tokens
+    loss_fct = nn.CrossEntropyLoss(reduction='none')
+    loss = loss_fct(shift_logits.permute(0, 2, 1), shift_labels).mean(dim=1)
+
+    return loss
+
+
+def rank0_print(msg):
+    if dist.get_rank() == 0:
+        print(msg)
+        
+
+def inference(config):
+    set_seed(config['seed'])
+
+    rank0_print(f"World size: {dist.get_world_size()}")
+
+    tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'], model_max_length=config['max_length'])
+    # llama has no pad token, set it to new token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+        
+    train_dataset, val_dataset = load_data_for_inference(config, tokenizer) 
+
+    num_processes = dist.get_world_size()
+    local_rank = dist.get_rank()
+
+    train_sampler = DistributedSampler(train_dataset, shuffle=False, drop_last=True, num_replicas=num_processes, rank=local_rank)
+    train_dataloader = DataLoader(
+        train_dataset,
+        collate_fn=DefaultDataCollator(),
+        batch_size=config["batch_size"],
+        sampler=train_sampler,
+        drop_last=True
+    )
+
+    val_sampler = DistributedSampler(val_dataset, shuffle=False, drop_last=True, num_replicas=num_processes, rank=local_rank)
+    val_dataloader = DataLoader(
+        val_dataset,
+        collate_fn=DefaultDataCollator(),
+        batch_size=config["batch_size"],
+        sampler=val_sampler,
+        drop_last=True
+    )
+
+
+    model = AutoModelForCausalLM.from_pretrained(config["model_name"], 
+                                                    trust_remote_code=True,
+                                                    torch_dtype=torch.bfloat16,
+                                                    ) 
+    model.to(f"cuda:{local_rank}")
+
+    with torch.no_grad():
+        train_outputs = {"loss": [], "embeddings": [], "index": []}
+        for batch in tqdm(train_dataloader, disable=local_rank != 0):
+            batch["input_ids"] = batch["input_ids"].to(f"cuda:{local_rank}")
+            batch["labels"] = batch["labels"].to(f"cuda:{local_rank}")
+            outputs = model(input_ids=batch["input_ids"], labels=batch["labels"], output_hidden_states=True)
+            loss = calc_cross_entropy_no_reduction(outputs.logits, batch["labels"])
+            train_outputs["loss"].extend(loss)
+
+            embeddings = outputs.hidden_states[-1]
+            batch_size = batch["input_ids"].shape[0]
+            sequence_lengths = []
+            # since we use mutiturn with multiple <|endoftext|>, we need to find the place where 
+            # <|endoftext|> is repeated
+            for item in batch["input_ids"]:
+                indices = torch.where(item == tokenizer.pad_token_id)[0]
+                found = False
+                for index in indices:
+                    # case where sequence is less than max length
+                    if torch.all(item[index:] == tokenizer.pad_token_id):
+                        sequence_lengths.append(index)
+                        found = True
+                        break
+                # case where sequence is >= max length
+                if not found:
+                    sequence_lengths.append(len(item) - 1)
+
+            sequence_lengths = torch.tensor(sequence_lengths)
+            pooled_logits = embeddings[torch.arange(batch_size, device=embeddings.device), sequence_lengths]
+
+            train_outputs["embeddings"].append(pooled_logits)
+            train_outputs["index"].extend(batch["index"].to(model.device))
+
+            torch.cuda.empty_cache()
+
+        train_outputs = nested_numpify(train_outputs)
+        # stack since they're 0-dim arrays
+        train_outputs["index"] = np.stack(train_outputs["index"])
+        train_outputs["loss"] = np.stack(train_outputs["loss"])
+        train_outputs["embeddings"] = np.concatenate(train_outputs["embeddings"])
+
+        df_train = Dataset.from_dict(train_outputs)
+        df_train = df_train.sort("index")
+        curr_idx = df_train["index"]
+
+        # compute mask in pyarrow since it's super fast
+        # ty @bmschmidt for showing me this!
+        table = train_dataset.data
+        mask = pc.is_in(table['index'], value_set=pa.array(curr_idx, pa.int32()))
+        filtered_table = table.filter(mask)
+        # convert from pyarrow to Dataset
+        filtered_train = Dataset.from_dict(filtered_table.to_pydict())
+
+        filtered_train = filtered_train.add_column("embeddings", df_train["embeddings"])
+        filtered_train = filtered_train.add_column("loss", df_train["loss"])
+        filtered_train = filtered_train.add_column("is_train", [True] * len(filtered_train))
+
+        filtered_train.to_json(f"inference/epoch_2_embeddings_train_shard_{local_rank}.jsonl", lines=True, orient="records", num_proc=64)
+
+        val_outputs = {"loss": [], "embeddings": [], "index": []}
+        for batch in tqdm(val_dataloader, disable=local_rank != 0):
+            batch["input_ids"] = batch["input_ids"].to(f"cuda:{local_rank}")
+            batch["labels"] = batch["labels"].to(f"cuda:{local_rank}")
+            outputs = model(input_ids=batch["input_ids"], labels=batch["labels"])
+            loss = calc_cross_entropy_no_reduction(outputs.logits, batch["labels"])
+            val_outputs["loss"].extend(loss)
+
+            logits = outputs.logits
+            batch_size = batch["input_ids"].shape[0]
+            sequence_lengths = []
+            # since we use mutiturn with multiple <|endoftext|>, we need to find the place where 
+            # <|endoftext|> is repeated
+            for item in batch["input_ids"]:
+                indices = torch.where(item == tokenizer.pad_token_id)[0]
+                found = False
+                for index in indices:
+                    if torch.all(item[index:] == tokenizer.pad_token_id):
+                        sequence_lengths.append(index)
+                        found = True
+                        break
+
+                # no match found
+                if not found:
+                    sequence_lengths.append(len(item) - 1)
+
+            sequence_lengths = torch.tensor(sequence_lengths)
+            pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+            val_outputs["embeddings"].append(pooled_logits)
+            val_outputs["index"].extend(batch["index"].to(model.device))
+
+            torch.cuda.empty_cache()
+
+        val_outputs = nested_numpify(val_outputs)
+        val_outputs["index"] = np.stack(val_outputs["index"])
+        val_outputs["loss"] = np.stack(val_outputs["loss"])
+        val_outputs["embeddings"] = np.concatenate(val_outputs["embeddings"])
+
+        df_val = Dataset.from_dict(val_outputs)
+        df_val = df_val.sort("index")
+        curr_idx = df_val["index"]
+
+        # compute mask in pyarrow since it's super fast
+        # ty @bmschmidt for showing me this!
+        table = val_dataset.data
+        mask = pc.is_in(table['index'], value_set=pa.array(curr_idx, pa.int32()))
+        filtered_table = table.filter(mask)
+        # convert from pyarrow to Dataset
+        filtered_val = Dataset.from_dict(filtered_table.to_pydict())
+
+        filtered_val = filtered_val.add_column("embeddings", df_val["embeddings"])
+        filtered_val = filtered_val.add_column("loss", df_val["loss"])
+        filtered_val = filtered_val.add_column("is_train", [False] * len(filtered_val))
+
+        filtered_val.to_json(f"inference/epoch_2_embeddings_val_shard_{local_rank}.jsonl", lines=True, orient="records", num_proc=64)
+    
+
+def main():
+    dist.init_process_group("nccl")
+    parser = ArgumentParser()
+    parser.add_argument("--config", type=str, default="config.yaml")
+
+    args = parser.parse_args()
+    config = read_config(args.config)
+
+    inference(config)
+
+
+if __name__ == "__main__":
+    # parse arguments by reading in a config
+    main()
+    

requirements.txt

@@ -9,4 +9,5 @@ peft
 nodelist-inflator
 deepspeed
 sentencepiece
-jsonlines
+jsonlines
+nomic

train.py

@@ -1,8 +1,7 @@
 import os
-from transformers import AutoModelForCausalLM, AutoTokenizer 
-from transformers.trainer_pt_utils import get_parameter_names
+from transformers import AutoModelForCausalLM, AutoTokenizer, get_scheduler, LlamaForCausalLM
 import torch
-import torch.nn as nn
+from torch.optim import AdamW
 from argparse import ArgumentParser
 from read import read_config
 from accelerate import Accelerator
@@ -11,7 +10,9 @@ from peft import get_peft_model, LoraConfig, TaskType
 from data import load_data
 from torchmetrics import MeanMetric
 from tqdm import tqdm
+import wandb
 
+torch.backends.cuda.matmul.allow_tf32 = True
 
 def format_metrics(metrics, split, prefix=""):
     log = f"[{split}]" + prefix
@@ -20,17 +21,12 @@ def format_metrics(metrics, split, prefix=""):
     return log
 
 
-def evaluate(config, model, val_dataloader):
+def evaluate(model, val_dataloader):
     model.eval()
-    val_loss = MeanMetric().to(model.device)
+    val_loss = MeanMetric(nan_strategy="error").to(model.device)
 
     with torch.no_grad():
-        for i, batch in enumerate(
-            tqdm(val_dataloader),
-        ):
-            if i == config["eval_steps"]:
-                break
-                
+        for batch in tqdm(val_dataloader):
             loss = model(**batch).loss
 
             loss_values = accelerator.gather_for_metrics({"loss": loss.detach()})
@@ -46,25 +42,20 @@ def train(accelerator, config):
     accelerator.print(config)
     accelerator.print(f"Using {accelerator.num_processes} GPUs")
 
-    tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'])
-    # llama has no pad token, set it to new token
+    tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_name'], model_max_length=config['max_length'])
+    # if no pad token, set it to eos
     if tokenizer.pad_token is None:
-        # these tokens are already in the vocab, just not mapped correctly
-        added_tokens = tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>", "pad_token": "<pad>"})
+        tokenizer.pad_token = tokenizer.eos_token
 
         
     with accelerator.main_process_first():
         train_dataloader, val_dataloader = load_data(config, tokenizer) 
-        
+
 
     checkpoint = config["gradient_checkpointing"]
     model = AutoModelForCausalLM.from_pretrained(config["model_name"], 
                                                     use_cache=False if checkpoint else True,
                                                     trust_remote_code=True) 
-
-    if added_tokens > 0:
-        model.resize_token_embeddings(len(tokenizer))
-    
     if checkpoint:
         model.gradient_checkpointing_enable()
 
@@ -77,7 +68,7 @@ def train(accelerator, config):
         model.print_trainable_parameters()
 
     optimizer_cls = (
-        torch.optim.AdamW
+        AdamW
         if accelerator.state.deepspeed_plugin is None
         or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
         else DummyOptim
@@ -85,11 +76,35 @@ def train(accelerator, config):
 
     # karpathy doesn't decay embeddding, maybe we should exclude
     # https://github.com/karpathy/minGPT/commit/bbbdac74fa9b2e55574d70056163ffbae42310c1#diff-2075fa9c224b395be5bda85544dd36572b59c76c54562819eadadbf268602834R157s
-    optimizer = optimizer_cls(model.parameters(), lr=config["lr"])
+    optimizer = optimizer_cls(model.parameters(), lr=config["lr"], weight_decay=config["weight_decay"])
 
-    # scheduler defined in Deepspeed config
-    scheduler = DummyScheduler(
-            optimizer,  warmup_num_steps=config["warmup_steps"],
+    if accelerator.state.deepspeed_plugin is not None:
+        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+
+    # decay to min_lr instead of 0
+    lr_ratio = config["min_lr"] / config["lr"]
+    accelerator.print(f"Len of train_dataloader: {len(train_dataloader)}")
+    total_num_steps = (len(train_dataloader) / gradient_accumulation_steps) * config["num_epochs"]
+    # instead of decaying to zero, decay to ratio of min_lr / lr
+    total_num_steps += int(total_num_steps * lr_ratio) + config["warmup_steps"]
+    accelerator.print(f"Total training steps: {total_num_steps}")
+
+    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        scheduler = get_scheduler(
+            name="cosine",
+            optimizer=optimizer,
+            num_warmup_steps=config["warmup_steps"] * accelerator.num_processes,
+            num_training_steps=total_num_steps * accelerator.num_processes,
+        )
+    else:
+        scheduler = DummyScheduler(
+            optimizer, total_num_steps=config["warmup_steps"], warmup_num_steps=config["warmup_steps"]
         )
 
     model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
@@ -108,21 +123,25 @@ def train(accelerator, config):
         accelerator.skip_first_batches(train_dataloader, resume_step)
         accelerator.print(f"Resuming from step {resume_step}")
 
-    train_loss = MeanMetric().to(model.device)
 
-    if accelerator.state.deepspeed_plugin is not None:
-        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
-            "gradient_accumulation_steps"
-        ]
+    # log gradients
+    if accelerator.is_main_process and config["wandb"]:
+        wandb.watch(model, log_freq=config["log_grads_every"], log="all")
 
     for epoch in range(config["num_epochs"]):
+        train_loss = MeanMetric(nan_strategy="error").to(model.device)
         for step, batch in enumerate(tqdm(train_dataloader)):
             model.train()
             outputs = model(**batch)
             loss = outputs.loss
-            loss = loss / gradient_accumulation_steps
 
+            # gather loss before backprop in case of gradient accumulation
+            loss_values = accelerator.gather_for_metrics({"loss": loss.detach().float()})
+            train_loss.update(loss_values["loss"])
+
+            loss = loss / gradient_accumulation_steps
             accelerator.backward(loss)
+            # get gradient norm of all params
 
             # log LR in case something weird happens 
             if step > 0 and step % (config["eval_every"] // 10) == 0:
@@ -135,14 +154,13 @@ def train(accelerator, config):
                 scheduler.step()
                 optimizer.zero_grad()
 
-            loss_values = accelerator.gather_for_metrics({"loss": loss.detach()})
-            train_loss.update(loss_values["loss"])
 
             if step > 0 and step % config["save_every"] == 0:
-                accelerator.save_state(f"{config['output_dir']}/step_{step}")
+                curr_step = step + epoch * len(train_dataloader)
+                accelerator.save_state(f"{config['output_dir']}/step_{curr_step}")
 
-            if step > 0 and step % config["eval_every"] == 0:
-                val_loss = evaluate(config, model, val_dataloader)
+            if step > 0 and (step % config["eval_every"] == 0 or step == len(train_dataloader) - 1):
+                val_loss = evaluate(model, val_dataloader)
 
                 log_train = {
                         "train_loss": train_loss.compute()
@@ -165,9 +183,20 @@ def train(accelerator, config):
         accelerator.print(f"Pushing to HF hub")
         accelerator.wait_for_everyone()
         unwrapped_model = accelerator.unwrap_model(model)
-        if accelerator.is_main_process:
-            unwrapped_model.push_to_hub(config["save_name"] + "_first_epoch", private=True)
+        try:
+            if accelerator.is_main_process:
+                unwrapped_model.push_to_hub(config["save_name"] + f"-epoch_{epoch}", private=True)
 
+        except Exception as e:
+            accelerator.print(e)
+            accelerator.print(f"Failed to push to hub")
+
+        unwrapped_model.save_pretrained(
+            f"{config['output_dir']}/-epoch_{epoch}",
+            is_main_process=accelerator.is_main_process,
+            save_function=accelerator.save,
+            state_dict=accelerator.get_state_dict(model),
+        )
             
     accelerator.wait_for_everyone()
     unwrapped_model = accelerator.unwrap_model(model)
@@ -178,9 +207,6 @@ def train(accelerator, config):
         state_dict=accelerator.get_state_dict(model),
     )
 
-    if accelerator.is_main_process:
-        unwrapped_model.push_to_hub(config["save_name"], private=True)
-
     accelerator.end_training()
 
     

worker_node_setup.sh

@@ -0,0 +1,6 @@
+#!/bin/sh
+HEAD_IP=$1
+
+sudo apt install -y nfs-common
+sudo mkdir -p ./data_multiplus
+sudo mount ${HEAD_IP}:${PWD}/data_multiplus ./data_multiplus