[WIP][Draft] PP + simplefsdp

stack-info: PR: #253, branch: IvanKobzarev/stack/11

Author: IvanKobzarev

autoparallel/_testing/models/dsv3.py

@@ -1556,6 +1556,7 @@ def forward(
             torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
         """
 
+        self.tok_embeddings = None
         h = self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens
 
         for layer in self.layers.values():
@@ -1630,6 +1631,7 @@ def init_weights(self, buffer_device: torch.device | None = None) -> None:
 
 def _init_weights_tok_embeddings(self: Union[DeepSeekV3Model, DeepSeekV3Stage0]):
     if self.tok_embeddings is not None:
+        torch.distributed.breakpoint()
         nn.init.normal_(self.tok_embeddings.weight)
 
 

autoparallel/_testing/models/llama3.py

@@ -178,6 +178,7 @@ def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Ten
     assert ndim > 1
     seqlen = x.shape[1]
     freqs_cis = freqs_cis[0:seqlen]
+    print(f"XXX FREQS_CIS.shape:{freqs_cis.shape} assert == {(seqlen, x.shape[-1])}")
     assert freqs_cis.shape == (seqlen, x.shape[-1])
     shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
     return freqs_cis.view(*shape)
@@ -297,7 +298,7 @@ def forward(
         xv = xv.view(bs, seqlen, -1, self.head_dim)
 
         # TODO: uncomment
-        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+        # xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
 
         # repeat k/v heads if n_kv_heads < n_heads
         keys = repeat_kv(xk, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)

autoparallel/apply_sharding.py

@@ -339,6 +339,9 @@ def apply_sharding_to_model(gm, sharding_placement, params_spec, buffers_spec):
     fqn_to_param = get_named_param_nodes(gm.graph)
     fqn_to_buffer = get_named_buffer_nodes(gm.graph)
 
+    # simple_fsdp_param_sharding
+    # simple_fsdp_mesh
+
     for fqn in params_spec:
         n = fqn_to_param[fqn]
         with unset_fake_temporarily():

autoparallel/init_weights.py

@@ -52,9 +52,9 @@ def getter(self) -> torch.nn.Parameter:
 
     def setter(self, value: Union[torch.Tensor, torch.nn.Parameter]) -> None:
         parallel_value = parallel_model.get_parameter(fqn)
-        assert isinstance(
-            parallel_value, DTensor
-        ), "Expected parallel_module params to be DTensors"
+        # assert isinstance(
+        #     parallel_value, DTensor
+        # ), "Expected parallel_module params to be DTensors"
         _copy_set_value_to_dtensor(fqn, parallel_value, value)
 
     return property(getter, setter)
@@ -66,9 +66,9 @@ def getter(self) -> torch.Tensor:
 
     def setter(self, value: torch.Tensor) -> None:
         parallel_value = parallel_model.get_buffer(fqn)
-        assert isinstance(
-            parallel_value, DTensor
-        ), "Expected parallel_module params to be DTensors"
+        # assert isinstance(
+        #     parallel_value, DTensor
+        # ), "Expected parallel_module params to be DTensors"
         _copy_set_value_to_dtensor(fqn, parallel_value, value)
 
     return property(getter, setter)