clean up

Author: wwwjn

scripts/checkpoint_conversion/convert_to_hf.py

@@ -44,7 +44,7 @@ def convert_to_hf(input_dir, output_dir, model_name, model_flavor, hf_assets_pat
     storage_writer = HuggingFaceStorageWriter(
         path=output_dir,
         save_distributed=True,
-        fqn_to_index_mapping=sd_adapter.fqn_to_index_mapping,
+        fqn_to_index_mapping=None,
         enable_consolidation=True,
         thread_count_consolidation=5,
     )

torchtitan/components/checkpoint.py

@@ -418,7 +418,7 @@ def dcp_load(
             )
 
             state_dict = self.sd_adapter.from_hf(hf_state_dict)
-            
+
             self.states[MODEL].load_state_dict(state_dict)
         else:
             dcp.load(state_dict, checkpoint_id=checkpoint_id)

torchtitan/models/deepseek_v3/__init__.py

@@ -46,7 +46,7 @@
             route_norm=True,
             score_before_experts=False,
         ),
-        q_lora_rank=0,
+        q_lora_rank=256,  # for test, original is 0
         kv_lora_rank=512,
         qk_nope_head_dim=128,
         qk_rope_head_dim=64,

torchtitan/models/deepseek_v3/hf_implementation.py

@@ -36,7 +36,6 @@ def parallelize_deepseekv3(
     job_config: JobConfig,
 ):
     world_mesh = parallel_dims.world_mesh
-    print(f"In parallelize_deepseekv3, world mesh is {world_mesh}")
     # TODO: TP currently cannot handle uneven seq_len because we set
     #       `use_local_output=True` to use plain Tensors for legacy reasons.
     #       Need to revisit this.

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -11,7 +11,7 @@
 from torch import nn
 
 from torchtitan.models.attention import build_attention
-from torchtitan.models.moe import FeedForward, MoE, print_tensor_stats
+from torchtitan.models.moe import FeedForward, MoE
 from torchtitan.protocols.train_spec import ModelProtocol
 
 from .args import DeepSeekV3ModelArgs
@@ -295,12 +295,9 @@ def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):
         Returns:
             torch.Tensor: Output tensor with the same shape as the input.
         """
-        print_tensor_stats(f"input of TransformerBlock {self.layer_id}: ", x)
         x = x + self.attention(self.attention_norm(x), freqs_cis)
         if self.moe_enabled:
-            x = self.ffn_norm(x)
-            print_tensor_stats(f"After ffn_norm : ", x)
-            x = x + self.moe(x)
+            x = x + self.moe(self.ffn_norm(x))
         else:
             x = x + self.feed_forward(self.ffn_norm(x))
         return x
@@ -388,11 +385,8 @@ def forward(
 
         h = self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens
 
-
-        token_inputs = h
         for layer in self.layers.values():
-            # reset before each layer
-            h = layer(token_inputs, self.freqs_cis)
+            h = layer(h, self.freqs_cis)
         h = self.norm(h) if self.norm is not None else h
         output = self.output(h) if self.output is not None else h
         return output