Update model.py

model.py

@@ -1,14 +1,419 @@
-from typing import Dict, Optional, Union
+from typing import Callable, Dict, Optional, Union, Tuple
 
 import copy
+import math
+import multiprocessing
+import os
+
 import torch
 import torch.nn as nn
 import transformers
 
-from cde.lib.dist import print0
-from cde.lib.tensor import mean_pool, mean_pool_3d, mean_pool_weighted, last_token_pool
 
-from cde.lib import load_embedder_and_tokenizer, ContextualModelConfig
+class ContextualModelConfig(transformers.configuration_utils.PretrainedConfig):
+    """We create a dummy configuration class that will just set properties
+    based on whatever kwargs we pass in.
+
+    When this class is initialized (see experiments.py) we pass in the
+    union of all data, model, and training args, all of which should
+    get saved to the config json.
+    """
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            try:
+                json.dumps(value)
+                setattr(self, key, value)
+            except TypeError:
+                # value was not JSON-serializable, skip
+                continue
+        super().__init__()
+
+
+def load_embedder_and_tokenizer(name: str) -> Tuple[
+        transformers.PreTrainedModel, 
+        transformers.PreTrainedTokenizer
+]:
+    if name.startswith("nomic") or (name == "bert-base-uncased"):
+        from cde.lib.nomic_bert import NomicBertModel
+        if name.endswith("--from-scratch"):
+            name = name.replace("--from-scratch", "")
+            config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
+            model = NomicBertModel._from_config(config)
+        else:
+            model = NomicBertModel.from_pretrained(
+                name, add_pooling_layer=False
+            )
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+    elif name in ["gtr-base", "gtr_base"]:
+        model = transformers.AutoModel.from_pretrained(
+            "sentence-transformers/gtr-t5-base"
+        ).encoder
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            "sentence-transformers/gtr-t5-base"
+        )
+    elif name == "pile-t5-base-encoder":
+        model = transformers.AutoModel.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        ).encoder
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+    elif name == "pile-t5-base-decoder":
+        model = transformers.AutoModel.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        ).decoder
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+    elif name.startswith("gpt2") or name.startswith("meta-llama") or ("Llama" in name):
+        model = transformers.AutoModelForCausalLM.from_pretrained(
+            name, 
+            # torch_dtype=torch.bfloat16,
+            attn_implementation="flash_attention_2",
+            low_cpu_mem_usage=True,
+            # device_map="auto",
+        )
+        model.padding_side = "right"
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.add_eos_token = True
+    else:
+        model = transformers.AutoModel.from_pretrained(name, trust_remote_code=True)
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+
+        # if use_bettertransformer:
+        #     from optimum.bettertransformer import BetterTransformer
+        #     model = BetterTransformer.transform(model)
+    return model, tokenizer
+def get_world_size() -> int:
+    try:
+        return torch.distributed.get_world_size()
+    except (RuntimeError, ValueError):
+        return 1
+
+
+def get_rank() -> int:
+    try:
+        return torch.distributed.get_rank()
+    except (RuntimeError, ValueError):
+        return 0
+    
+def gather(t: torch.Tensor) -> torch.Tensor:
+    # torch.distributed.nn.all_gather scales by world size since the reduce op is SUM
+    # https://github.com/pytorch/pytorch/issues/58005
+    # only should use torch.distributed.nn.all_gather if we implement a `local_loss`
+    # like: https://github.com/mlfoundations/open_clip/issues/616
+    world_size = get_world_size()
+    if world_size == 1:
+        return t
+
+    if t.ndim == 0:
+        t = t.unsqueeze(0)
+
+    gathered = [torch.empty_like(t) for _ in range(world_size)]
+    torch.distributed.all_gather(gathered, t)
+    gathered[get_rank()] = t
+    return torch.cat(gathered, dim=0)
+
+
+def gather_sum(t: torch.Tensor) -> torch.Tensor:
+    # torch.distributed.nn.all_gather scales by world size since the reduce op is SUM
+    # https://github.com/pytorch/pytorch/issues/58005
+    # only should use torch.distributed.nn.all_gather if we implement a `local_loss`
+    # like: https://github.com/mlfoundations/open_clip/issues/616
+    world_size = get_world_size()
+    if world_size == 1:
+        return t
+
+    if t.ndim == 0:
+        t = t.unsqueeze(0)
+
+    gathered = [torch.empty_like(t) for _ in range(world_size)]
+    torch.distributed.all_gather(gathered, t)
+    gathered = torch.stack(gathered, dim=0)
+    return gathered.sum(dim=0) # Sum across workers
+
+
+def get_num_proc() -> int:
+    world_size: int = get_world_size()
+    try:
+        # os.sched_getaffinity respects schedulers, unlike cpu_count(), but it's only available
+        # on some Unix platforms, so we support both!
+        return len(os.sched_getaffinity(0)) // world_size  # type: ignore[attr-defined]
+    except AttributeError:
+        return multiprocessing.cpu_count() // world_size
+
+
+def torch_main_worker_finish_first(func: Callable):
+    def wrapper(*args, **kwargs):
+        # Get local rank (need to support non-DDP).
+        try:
+            local_rank = torch.distributed.get_rank()
+            ddp_enabled = True
+        except (RuntimeError, ValueError):
+            local_rank = -1
+            ddp_enabled = False
+        is_main_worker = local_rank <= 0
+        # Run on main worker first.
+        if is_main_worker:
+            result = func(*args, **kwargs)
+        # Then everyone waits.
+        if ddp_enabled:
+            torch.distributed.barrier()
+        # Run on other workers now.
+        if not is_main_worker:
+            result = func(*args, **kwargs)
+        # Now everyone waits again.
+        if ddp_enabled:
+            torch.distributed.barrier()
+        return result
+
+    return wrapper
+
+
+def print0(*args, **kwargs) -> None:
+    if get_rank() == 0:
+        print(*args, **kwargs)
+
+
+def verify_ddp_weights_equal(model: torch.nn.Module, atol: float = 1e-5) -> None:
+    if hasattr(model, "module"):
+        model = model.module
+    
+    world_size = get_world_size()
+
+    if world_size > 8:
+        print0(f"[verify_ddp_weights_equal] Skipping with world_size={world_size} ⚠️")
+        return
+
+    for name, param in model.named_parameters():
+        if param is None: continue
+        if param.grad is None: 
+            print0(f"[verify_ddp_weights_equal] Skipping param [{name}] with no grad")
+            continue
+        gathered_param = gather(param).reshape((world_size, -1))
+        absolute_diffs = (gathered_param[None, 0, :] - gathered_param).abs()
+        rank_params_eq = (absolute_diffs < atol).all()
+        assert rank_params_eq, f"❌ param [{name}] not equal - got max_absolute_diff={absolute_diffs.max()}"
+        ###################################################################################################################
+        gathered_param_grad = gather(param.grad).reshape((world_size, -1))
+        absolute_grad_diffs = (gathered_param_grad[None, 0, :] - gathered_param_grad).abs()
+        rank_grad_params_eq = (absolute_grad_diffs < atol).all()
+        assert rank_grad_params_eq, f"❌ param [{name}] grad not equal - got max_absolute_diff={absolute_grad_diffs.max()}"
+        ###################################################################################################################
+        
+    
+    print0("[verify_ddp_weights_equal] Verified DDP parameter correctness ✅")
+    
+
+
+def mean_pool_3d(
+    hidden_states: torch.Tensor, attention_mask: torch.Tensor
+) -> torch.Tensor:
+    B, T, S, D = hidden_states.shape
+    unmasked_outputs = hidden_states * attention_mask[..., None]
+    pooled_outputs = unmasked_outputs.sum(dim=2) / (attention_mask.sum(dim=2)[..., None] + 1e-9)
+
+    # fix for gradient flow: fill empty rows with the mean of the rest of the sequence
+    sequence_means = (
+        hidden_states.reshape((B, S * T, D))
+            .mean(dim=1, keepdim=True)
+            .expand(-1, T, -1)
+    )
+    pooled_outputs = pooled_outputs.where(
+        (attention_mask.sum(dim=2)[..., None] > 0), 
+        sequence_means
+    )
+    assert pooled_outputs.shape == (B, T, D)
+
+    return pooled_outputs
+
+def mean_pool(
+    hidden_states: torch.Tensor, attention_mask: torch.Tensor
+) -> torch.Tensor:
+    B, _S, D = hidden_states.shape
+    unmasked_outputs = hidden_states * attention_mask[..., None]
+    pooled_outputs = unmasked_outputs.sum(dim=1) / (attention_mask.sum(dim=1)[:, None] + 1e-20)
+    
+    assert pooled_outputs.shape == (B, D)
+    return pooled_outputs
+
+
+def mean_pool_weighted(
+    hidden_states: torch.Tensor, attention_mask: torch.Tensor
+) -> torch.Tensor:
+    B, _S, D = hidden_states.shape
+    attention_mask *= attention_mask.cumsum(dim=1) # [0,1,1,1,0,0] -> [0,1,2,3,0,0]
+    s = torch.sum(hidden_states * attention_mask.unsqueeze(-1).float(), dim=1)
+    d = attention_mask.sum(dim=1, keepdim=True).float()
+    return s / d
+
+
+def slice_sparse_tensor_rows(t: torch.sparse.Tensor, min_row: int, max_row: int) -> torch.sparse.Tensor:
+    assert min_row < max_row, f"can't slice from row {min_row} to {max_row}"
+    t = t.coalesce()
+    row_idxs = t.indices()[0]
+    index_mask = (min_row <= row_idxs) & (row_idxs < max_row)
+
+    num_rows = (max_row - min_row)
+    num_cols = t.shape[1]
+
+    idxs = t.indices()[:, index_mask]
+    vals = t.values()[index_mask]
+    return torch.sparse_coo_tensor(idxs, vals, size=(num_rows, num_cols)).coalesce()
+
+
+def slice_tensor_rows(t: torch.Tensor, min_row: int, max_row: int) -> torch.Tensor:
+    if t.is_sparse:
+        return slice_sparse_tensor_rows(t=t, min_row=min_row, max_row=max_row)
+    else:
+        return t[min_row:max_row]
+
+
+@torch.no_grad
+def maxsim(
+    X: torch.Tensor, y: torch.Tensor, 
+    maximize: bool, chunk_size: int = 8_000,
+    debug_mem_usage: bool = False) -> torch.Tensor:
+    device = X.device
+    n_samples = X.shape[0]
+
+    max_sim_v = torch.zeros(n_samples, device=device, dtype=X.dtype)
+    max_sim_i = torch.zeros(n_samples, device=device, dtype=torch.int64)
+
+    # TODO: Implement faster max (without going to dense tensors).
+    # TODO: Use multiple GPUs.
+    rank = get_rank()
+    world_size = get_world_size()
+
+    worker_worklist_size = int(math.ceil(n_samples / world_size))
+    splits_start_idx = worker_worklist_size * rank
+    splits_end_idx = worker_worklist_size * (rank + 1)
+
+    for i in range(splits_start_idx, splits_end_idx, chunk_size):
+        start, end = i, min(i + chunk_size, n_samples)
+        sub_x = slice_tensor_rows(X, start, end)
+        if debug_mem_usage: print(f"[maxsim] step {i} cuda mem free/total = {torch.cuda.mem_get_info()}")
+        if debug_mem_usage: print("[maxsim] sub_x.shape:", sub_x.shape, "//", "y.shape:", y.shape)
+        sub_sim = sub_x @ y # TODO – Implement sparse max here to save mem!
+        sub_sim = sub_sim
+        if maximize:
+            sub_max_sim_v, sub_max_sim_i = sub_sim.to_dense().max(dim=-1)
+        else:
+            sub_max_sim_v, sub_max_sim_i = sub_sim.to_dense().min(dim=-1)
+        del sub_sim
+        del sub_x
+        torch.cuda.empty_cache() # needs to happen after maxsim for some reason.
+        max_sim_v[start: end] = sub_max_sim_v
+        max_sim_i[start: end] = sub_max_sim_i
+    
+    # gather
+    max_sim_v = gather_sum(max_sim_v)
+    max_sim_i = gather_sum(max_sim_i)
+    k = y.shape[1]
+
+    assert max_sim_v.shape == (n_samples,)
+    assert max_sim_i.shape == (n_samples,)
+    assert max_sim_i.min() >= 0
+    assert max_sim_i.max() <= k
+
+    return max_sim_v, max_sim_i
+
+
+def forward_batched(
+        model: torch.nn.Module,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        batch_size: int,
+        dataset_input_ids: Optional[torch.Tensor] = None,
+        dataset_attention_mask: Optional[torch.Tensor] = None,
+        **second_stage_model_kwargs,
+) -> torch.Tensor:
+    if hasattr(model, "module"):
+        model = model.module
+    
+    if hasattr(model, "first_stage_model"):
+        # Support pooling over 3D dataset_input_ids inputs.
+        if len(dataset_input_ids.shape) == 2:
+            dataset_input_ids = dataset_input_ids[None]
+            dataset_attention_mask = dataset_attention_mask[None]
+
+        dataset_embeddings = []
+        for j in range(len(dataset_input_ids)):
+            i = 0
+            dataset_embeddings_batch = []
+            while i < dataset_input_ids.shape[1]:
+                dataset_embeddings_batch.append(
+                    model.first_stage_model(
+                        input_ids=dataset_input_ids[j][i:i+batch_size],
+                        attention_mask=dataset_attention_mask[j][i:i+batch_size],
+                    )
+                )
+                i += batch_size
+            dataset_embeddings.append(
+                torch.cat(dataset_embeddings_batch, dim=0)
+            )
+       
+        # Automatically pool over 3D dataset_input_ids.
+        dataset_embeddings = torch.stack(dataset_embeddings, dim=0).mean(dim=0)
+
+        j = 0
+        outputs = []
+        while j < len(input_ids):
+            outputs.append(
+                model.second_stage_model(
+                    input_ids=input_ids[j:j+batch_size],
+                    attention_mask=attention_mask[j:j+batch_size],
+                    dataset_embeddings=dataset_embeddings,
+                    **second_stage_model_kwargs,
+                )
+            )
+            j += batch_size
+        return torch.cat(outputs, dim=0)
+
+    else:
+        i = 0
+        outputs = []
+        while i < len(input_ids):
+            # breakpoint()
+            outputs.append(
+                model(
+                    input_ids=input_ids[i:i+batch_size],
+                    attention_mask=attention_mask[i:i+batch_size],
+                    **second_stage_model_kwargs,
+                )
+            )
+            i += batch_size
+        return torch.cat(outputs, dim=0)
+
+
+def last_token_pool(hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    # https://github.com/ContextualAI/gritlm/blob/main/gritlm/gritlm.py#L190
+    b, n, d = hidden_state.size()
+    # Get the last `1` in the attention mask of each item
+    # Often it is just `gather_indices = torch.argmin(attention_mask, 1, keepdim=False) - 1`
+    # except when 1) There's all 1's 2) There's 0's before the 1's
+    reversed_mask = torch.flip(attention_mask, dims=(1,))
+    argmax_reverse = torch.argmax(reversed_mask, dim=1, keepdim=False)
+    gather_indices = attention_mask.size(1) - argmax_reverse - 1
+    # If there are empty sequences, where the index would become -1 it will crash so set them to 0
+    gather_indices = torch.clamp(gather_indices, min=0)
+    # Turn indices from shape [b] -> [b, 1, d]
+    gather_indices = gather_indices.unsqueeze(-1).repeat(1, d)
+    gather_indices = gather_indices.unsqueeze(1)
+    assert gather_indices.shape == (b, 1, d)
+    # Gather along the seq len: [b, n, d] -> [b, d]
+    # Actually no need for the attention mask as we gather the last token where attn_mask=1 but
+    # as some indices (which shouldn't be attended to) may be 0 due to clamp, use mask to ignore them again
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand((b, n, d)).float()
+    return torch.gather(hidden_state * input_mask_expanded, 1, gather_indices).squeeze(dim=1)
+
+def print0(*args, **kwargs) -> None:
+    if get_rank() == 0:
+        print(*args, **kwargs)
 
 
 def limit_layers(model: transformers.PreTrainedModel, n_layers: int) -> None:
@@ -25,6 +430,7 @@ def limit_layers(model: transformers.PreTrainedModel, n_layers: int) -> None:
             model.encoder.layer = model.encoder.layer[:n_layers]
     else:
         raise RuntimeError(f"unknown how to limit layers of model {type(model)}")
+    
 
 
 def disable_dropout(model: torch.nn.Module):
@@ -413,7 +819,7 @@ class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelM
                 if hasattr(module, "rotary_emb_dim"):
                     module.rotary_start_pos = rotary_start_pos
                     rotary_disabled += 1
-            print0(f"modified {rotary_disabled} rotary modules – set rotary_start_pos to {rotary_start_pos}")
+            print0(f"modified {rotary_disabled} rotary modules – set rotary_start_pos to {rotary_start_pos}")
     
     def forward(
             self, 
@@ -580,7 +986,7 @@ class DatasetTransformer(transformers.PreTrainedModel):
             output_hidden_states: bool = False,
         ) -> torch.Tensor:
         """
-        input_ids (long torch.Tensor) – ids of input tokens
+        input_ids (long torch.Tensor) – ids of input tokens
         attention_mask (bool torch.Tensor)
         """
         dataset_embeddings = self.first_stage_model(