Merge branch 'dev' into main

ndif-team · Jan 31, 2025 · 80ac87f · 80ac87f
2 parents 824c0a8 + eccb2fb
commit 80ac87f
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 1,238 deletions.
diff --git a/src/nnsight/__init__.py b/src/nnsight/__init__.py
@@ -79,6 +79,22 @@ def inner(cls, fake_mode, elem, device, constant=None):
     Patch(FakeTensor, fake_tensor_new_wrapper(FakeTensor.__new__), "__new__")
 )
 
+def autoamp_init_wrapper(fn):
+
+    @wraps(fn)
+    def inner(self, device_type, dtype=None, **kwargs):
+
+        if device_type == "meta":
+            dtype = torch.get_autocast_cpu_dtype()
+
+        return fn(self, device_type, dtype, **kwargs)
+
+    return inner
+
+DEFAULT_PATCHER.add(
+    Patch(torch.autocast, autoamp_init_wrapper(torch.autocast.__init__), "__init__")
+)
+
 DEFAULT_PATCHER.__enter__()
 
 from .intervention.contexts import GlobalInterventionTracingContext

diff --git a/src/nnsight/models/LanguageModel.py b/src/nnsight/models/LanguageModel.py