Merge pull request #117 from EleutherAI/convert-dtype

Use dtype utils

Author: luciaquirke

bergson/collector/dist_preconditioners_gradient_collector.py

@@ -18,7 +18,7 @@
 )
 from bergson.process_preconditioners import process_preconditioners
 from bergson.score.scorer import Scorer
-from bergson.utils.utils import assert_type
+from bergson.utils.utils import assert_type, get_gradient_dtype
 
 
 @dataclass(kw_only=True)
@@ -98,18 +98,14 @@ def setup(self) -> None:
         self.owned_modules: set[str] = set()
         self.module_to_rank: dict[str, int] = {}
 
-        # TODO: handle more elegantly?
-        self.save_dtype = (
-            torch.float32 if self.model.dtype == torch.float32 else torch.float16
-        )
-
+        self.save_dtype = get_gradient_dtype(self.model)
         self.lo = torch.finfo(self.save_dtype).min
         self.hi = torch.finfo(self.save_dtype).max
 
         self.per_doc_losses = torch.full(
             (len(self.data),),
             device=self.model.device,
-            dtype=self.save_dtype,
+            dtype=torch.float32,
             fill_value=0.0,
         )
 
@@ -298,11 +294,7 @@ def teardown(self):
             self.data = self.data.add_column(
                 "loss",
                 self.per_doc_losses.cpu().numpy(),
-                feature=Value(
-                    "float16"
-                    if self.save_dtype == torch.float16
-                    else "float32"  # TODO: This is not robust
-                ),
+                feature=Value("float32"),
                 new_fingerprint="loss",
             )
 

bergson/collector/gradient_collectors.py

@@ -19,7 +19,7 @@
 )
 from bergson.process_preconditioners import process_preconditioners
 from bergson.score.scorer import Scorer
-from bergson.utils.utils import assert_type
+from bergson.utils.utils import assert_type, get_gradient_dtype
 
 
 @dataclass(kw_only=True)
@@ -93,18 +93,14 @@ def setup(self) -> None:
                 "consider disabling bias inclusion for now."
             )
 
-        # TODO: handle more elegantly?
-        self.save_dtype = (
-            torch.float32 if self.model.dtype == torch.float32 else torch.float16
-        )
-
+        self.save_dtype = get_gradient_dtype(self.model)
         self.lo = torch.finfo(self.save_dtype).min
         self.hi = torch.finfo(self.save_dtype).max
 
         self.per_doc_losses = torch.full(
             (len(self.data),),
             device=self.model.device,
-            dtype=self.save_dtype,
+            dtype=torch.float32,
             fill_value=0.0,
         )
 
@@ -263,11 +259,7 @@ def teardown(self):
                 self.data = self.data.add_column(
                     "loss",
                     self.per_doc_losses.cpu().numpy(),
-                    feature=Value(
-                        "float16"
-                        if self.save_dtype == torch.float16
-                        else "float32"  # TODO: This is not robust
-                    ),
+                    feature=Value("float32"),
                     new_fingerprint="loss",
                 )
 
@@ -302,11 +294,7 @@ class TraceCollector(HookCollectorBase):
     """Dtype for stored gradients."""
 
     def setup(self) -> None:
-        # TODO: handle more elegantly?
-        self.save_dtype = (
-            torch.float32 if self.model.dtype == torch.float32 else torch.float16
-        )
-
+        self.save_dtype = get_gradient_dtype(self.model)
         self.lo = torch.finfo(self.save_dtype).min
         self.hi = torch.finfo(self.save_dtype).max
 

bergson/config.py

@@ -282,6 +282,10 @@ class ScoreConfig:
     batch_size: int = 1024
     """Batch size for processing the query dataset."""
 
+    precision: Literal["auto", "bf16", "fp16", "fp32"] = "auto"
+    """Precision (dtype) to convert the query and index gradients to before
+    computing the scores. If "auto", the model's gradient dtype is used."""
+
 
 @dataclass
 class ReduceConfig:

bergson/data.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import Any, Sequence, cast, overload
 
+import ml_dtypes  # noqa: F401  # registers bfloat16 dtype with numpy
 import numpy as np
 import pyarrow as pa
 import torch
@@ -20,7 +21,12 @@
 from numpy.typing import DTypeLike
 
 from .config import DataConfig, ReduceConfig
-from .utils.utils import assert_type, simple_parse_args_string
+from .utils.utils import (
+    assert_type,
+    convert_dtype_to_np,
+    simple_parse_args_string,
+    tensor_to_numpy,
+)
 
 
 def ceildiv(a: int, b: int) -> int:
@@ -202,7 +208,7 @@ def create_index(
                     "num_grads": num_grads,
                     "dtype": struct_dtype,
                     "grad_sizes": grad_sizes,
-                    "base_dtype": np.dtype(dtype).str,
+                    "base_dtype": np.dtype(dtype).name,
                 },
                 f,
                 indent=2,
@@ -367,17 +373,16 @@ def __init__(
         self.rank = dist.get_rank() if dist.is_initialized() else 0
         if reduce_cfg is not None:
             num_grads = 1
+            np_dtype = np.float32
             self.in_memory_grad_buffer = torch.zeros(
                 (num_grads, sum(self.grad_sizes.values())),
                 dtype=torch.float32,
                 device=f"cuda:{self.rank}",
             )
-            np_dtype = np.float32
         else:
             num_grads = self.num_items
+            np_dtype = convert_dtype_to_np(dtype)
             self.in_memory_grad_buffer = None
-            # TODO: Handle this more elegantly
-            np_dtype = np.float32 if dtype == torch.float32 else np.float16
 
         self.grad_buffer = create_index(
             path,
@@ -423,7 +428,7 @@ def __call__(self, indices: list[int], mod_grads: dict[str, torch.Tensor]):
             for module_name in self.grad_sizes.keys():
                 self.grad_buffer[
                     indices, offset : offset + mod_grads[module_name].shape[1]
-                ] = mod_grads[module_name].numpy()
+                ] = tensor_to_numpy(mod_grads[module_name])
                 offset += mod_grads[module_name].shape[1]
 
     def flush(self):
@@ -447,7 +452,7 @@ def dist_reduce(self):
 
         rank = dist.get_rank() if dist.is_initialized() else 0
         if rank == 0:
-            self.grad_buffer[:] = self.in_memory_grad_buffer.numpy().astype(
+            self.grad_buffer[:] = tensor_to_numpy(self.in_memory_grad_buffer).astype(
                 self.grad_buffer.dtype
             )
 

bergson/huggingface.py

@@ -9,7 +9,6 @@
 import torch
 import torch.distributed as dist
 from datasets import Dataset
-from numpy.typing import DTypeLike
 from peft import PeftModel
 from torch import Tensor
 from torch.utils.data import DataLoader
@@ -22,6 +21,7 @@
 from bergson.data import create_index
 from bergson.gradients import AdafactorNormalizer, AdamNormalizer
 from bergson.utils.peft import detect_peft_modules
+from bergson.utils.utils import convert_dtype_to_torch
 
 
 class GradientCollectorCallback(TrainerCallback):
@@ -34,7 +34,7 @@ def __init__(
         attention_cfgs: dict[str, AttentionConfig] = {},
         projection_dim: int = 16,
         include_bias: bool = False,
-        dtype: DTypeLike = np.float16,
+        dtype: np.dtype = np.dtype(np.float16),
         accumulate_grads: bool = False,
         use_optimizer_state: bool = True,
         track_order: bool = False,
@@ -77,8 +77,7 @@ def __init__(
         self.mod_grads = {}
         self.batch_indices: Tensor | None = None
 
-        # TODO: Handle this more elegantly
-        self.torch_dtype = torch.float32 if self.dtype == np.float32 else torch.float16
+        self.torch_dtype = convert_dtype_to_torch(self.dtype)
 
     def write_grads(self, grad_buffer: np.memmap):
         torch.cuda.synchronize()

bergson/normalizer/fit_normalizers.py

@@ -19,7 +19,7 @@
     Normalizer,
 )
 from bergson.process_preconditioners import process_preconditioners
-from bergson.utils.utils import assert_type
+from bergson.utils.utils import assert_type, get_gradient_dtype
 
 
 @dataclass(kw_only=True)
@@ -123,11 +123,7 @@ def setup(self) -> None:
                 "consider disabling bias inclusion for now."
             )
 
-        # TODO: handle more elegantly?
-        self.save_dtype = (
-            torch.float32 if self.model.dtype == torch.float32 else torch.float16
-        )
-
+        self.save_dtype = get_gradient_dtype(self.model)
         self.lo = torch.finfo(self.save_dtype).min
         self.hi = torch.finfo(self.save_dtype).max
 

bergson/score/score.py

@@ -17,7 +17,11 @@
 from bergson.distributed import launch_distributed_run
 from bergson.gradients import GradientProcessor
 from bergson.score.scorer import Scorer
-from bergson.utils.utils import assert_type
+from bergson.utils.utils import (
+    assert_type,
+    convert_precision_to_torch,
+    get_gradient_dtype,
+)
 from bergson.utils.worker_utils import (
     create_processor,
     setup_data_pipeline,
@@ -277,6 +281,12 @@ def score_worker(
         "attention_cfgs": attention_cfgs,
     }
 
+    score_dtype = (
+        convert_precision_to_torch(score_cfg.precision)
+        if score_cfg.precision != "auto"
+        else get_gradient_dtype(model)
+    )
+
     if isinstance(ds, Dataset):
         kwargs["batches"] = allocate_batches(ds["length"], index_cfg.token_batch_size)
         kwargs["scorer"] = Scorer(
@@ -285,7 +295,7 @@ def score_worker(
             query_grads,
             score_cfg,
             device=torch.device(f"cuda:{rank}"),
-            dtype=torch.float32 if model.dtype == torch.float32 else torch.float16,
+            dtype=score_dtype,
         )
 
         collect_gradients(**kwargs)
@@ -310,7 +320,7 @@ def flush(kwargs):
                 query_grads,
                 score_cfg,
                 torch.device(f"cuda:{rank}"),
-                model.dtype if model.dtype != "auto" else torch.float32,
+                score_dtype,
             )
 
             collect_gradients(**kwargs)

bergson/score/scorer.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from typing import Callable
 
 import torch
 
@@ -8,14 +7,8 @@
 
 
 class Scorer:
-    scorer_callback: Callable
-
-    num_scores: int
-
     writer: ScoreWriter
 
-    device: torch.device
-
     def __init__(
         self,
         path: Path,
@@ -29,10 +22,14 @@ def __init__(
         self.dtype = dtype
         self.num_items = num_items
 
-        self.scorer_callback = self.build_scorer_callback(
-            query_grads,
-            score_cfg,
+        self.query_tensor = torch.cat(
+            [
+                query_grads[m].to(device=self.device, dtype=self.dtype)
+                for m in score_cfg.modules
+            ],
+            dim=1,
         )
+        self.score_cfg = score_cfg
 
         num_scores = len(query_grads[score_cfg.modules[0]])
 
@@ -47,37 +44,22 @@ def __call__(
         indices: list[int],
         mod_grads: dict[str, torch.Tensor],
     ):
-        first_grad = next(iter(mod_grads.values()))
-        if first_grad.dtype != self.dtype:
+        # Convert the gradients to the scoring dtype
+        if next(iter(mod_grads.values())).dtype != self.dtype:
             mod_grads = {name: grad.to(self.dtype) for name, grad in mod_grads.items()}
 
-        scores = self.scorer_callback(mod_grads)
-        self.writer(indices, scores)
-
-    def build_scorer_callback(
-        self,
-        query_grads: dict[str, torch.Tensor],
-        score_cfg: ScoreConfig,
-    ) -> Callable:
-        """Unified scorer builder for all scorer types."""
-        query_tensor = torch.cat(
-            [
-                query_grads[m].to(device=self.device, dtype=self.dtype)
-                for m in score_cfg.modules
-            ],
-            dim=1,
-        )
+        scores = self.score(mod_grads)
 
-        @torch.inference_mode()
-        def callback(mod_grads: dict[str, torch.Tensor]):
-            grads = torch.cat([mod_grads[m] for m in score_cfg.modules], dim=1)
-            if score_cfg.unit_normalize:
-                grads /= grads.norm(dim=1, keepdim=True)
+        self.writer(indices, scores)
 
-            if score_cfg.score == "nearest":
-                all_scores = grads @ query_tensor.T
-                return all_scores.max(dim=-1).values
+    @torch.inference_mode()
+    def score(self, mod_grads: dict[str, torch.Tensor]):
+        grads = torch.cat([mod_grads[m] for m in self.score_cfg.modules], dim=1)
+        if self.score_cfg.unit_normalize:
+            grads /= grads.norm(dim=1, keepdim=True)
 
-            return grads @ query_tensor.T
+        if self.score_cfg.score == "nearest":
+            all_scores = grads @ self.query_tensor.T
+            return all_scores.max(dim=-1).values
 
-        return callback
+        return grads @ self.query_tensor.T