Lifting check_unitialized.

server/text_generation_server/models/bloom.py

@@ -238,14 +238,7 @@ class BLOOMSharded(BLOOM):
                     if name == "word_embeddings.weight":
                         model.lm_head._parameters["weight"] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
 
     def forward(
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None

server/text_generation_server/models/flash_llama.py

@@ -139,15 +139,7 @@ class FlashLlama(FlashCausalLM):
 
                 del value
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
-
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)
 
@@ -315,14 +307,7 @@ class FlashLlamaSharded(FlashLlama):
                     else:
                         module._buffers[param_name] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
 
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)

server/text_generation_server/models/flash_neox.py

@@ -152,13 +152,5 @@ class FlashNeoXSharded(FlashNeoX):
                     else:
                         module._buffers[param_name] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
-
         model.post_load_weights(quantize)

server/text_generation_server/models/flash_santacoder.py

@@ -376,17 +376,7 @@ class FlashSantacoderSharded(FlashSantacoder):
                     else:
                         module._buffers[param_name] = tensor
 
-
         model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight)
 
-        uninitialized_parameters = []
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
-
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)

server/text_generation_server/models/galactica.py

@@ -365,14 +365,7 @@ class GalacticaSharded(Galactica):
                     if name == "model.decoder.embed_tokens.weight":
                         model.lm_head._parameters["weight"] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
 
     def forward(
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None

server/text_generation_server/models/gpt_neox.py

@@ -215,14 +215,7 @@ class GPTNeoxSharded(CausalLM):
                     else:
                         module._buffers[param_name] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
 
     def forward(
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None

server/text_generation_server/models/model.py

@@ -99,3 +99,13 @@ class Model(ABC):
             return token_text, None, None
         else:
             return "", offset, token_offset
+
+    def check_initialized(self):
+        uninitialized_parameters = []
+        for n, p in self.named_parameters():
+            if p.data.device == torch.device("meta"):
+                uninitialized_parameters.append(n)
+        if uninitialized_parameters:
+            raise RuntimeError(
+                f"found uninitialized parameters in model: {uninitialized_parameters}"
+            )

server/text_generation_server/models/opt.py

@@ -212,14 +212,7 @@ class OPTSharded(OPT):
                     if name == "model.decoder.embed_tokens.weight":
                         model.lm_head._parameters["weight"] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
 
     def forward(
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None

server/text_generation_server/models/t5.py

@@ -222,14 +222,7 @@ class T5Sharded(Seq2SeqLM):
                     else:
                         module._buffers[param_name] = tensor
 
-        uninitialized_parameters = []
+        model.check_initialized()
-        for n, p in model.named_parameters():
-            if p.data.device == torch.device("meta"):
-                uninitialized_parameters.append(n)
-        if uninitialized_parameters:
-            raise RuntimeError(
-                f"found uninitialized parameters in model: {uninitialized_parameters}"
-            )
 
     def forward(
         self,