LLM-foundry update June 16, 2023 22:55:57

custom_embedding.py

@@ -3,10 +3,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 
-
 class SharedEmbedding(nn.Embedding):
 
-    def forward(self, input: Tensor, unembed: bool = False) -> Tensor:
+    def forward(self, input: Tensor, unembed: bool=False) -> Tensor:
         if unembed:
             return F.linear(input, self.weight)
         return super().forward(input)

modeling_mpt.py

@@ -40,6 +40,11 @@ class MPTModel(MPTPreTrainedModel):
         self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
         self.alibi = config.attn_config['alibi']
         self.alibi_bias_max = config.attn_config['alibi_bias_max']
+        if config.init_device == 'mixed':
+            if dist.get_local_rank() == 0:
+                config.init_device = 'cpu'
+            else:
+                config.init_device = 'meta'
         if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():
             norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
             raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
@@ -47,7 +52,7 @@ class MPTModel(MPTPreTrainedModel):
         self.embedding_fraction = config.embedding_fraction
         self.wte = SharedEmbedding(config.vocab_size, config.d_model, device=config.init_device)
         if not self.alibi:
-            self.wpe = nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
+            self.wpe = torch.nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
         self.emb_drop = nn.Dropout(config.emb_pdrop)
         self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
         self.norm_f = norm_class(config.d_model, device=config.init_device)
@@ -221,6 +226,11 @@ class MPTForCausalLM(MPTPreTrainedModel):
         if not config.tie_word_embeddings:
             raise ValueError('MPTForCausalLM only supports tied word embeddings')
         self.transformer = MPTModel(config)
+        for child in self.transformer.children():
+            if isinstance(child, torch.nn.ModuleList):
+                continue
+            if isinstance(child, torch.nn.Module):
+                child._fsdp_wrap = True
         self.logit_scale = None
         if config.logit_scale is not None:
             logit_scale = config.logit_scale