further 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=None,
+        fqn_to_index_mapping=sd_adapter.fqn_to_index_mapping,
         enable_consolidation=True,
         thread_count_consolidation=5,
     )

torchtitan/components/checkpoint.py

@@ -418,7 +418,6 @@ 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=256,  # for test, original is 0
+        q_lora_rank=0,
         kv_lora_rank=512,
         qk_nope_head_dim=128,
         qk_rope_head_dim=64,
@@ -135,7 +135,7 @@
         dim=7168,
         inter_dim=18432,
         moe_inter_dim=2048,
-        n_layers=4,
+        n_layers=61,
         n_dense_layers=3,
         n_heads=128,
         moe_args=MoEArgs(

torchtitan/models/deepseek_v3/model/state_dict_adapter.py

@@ -378,8 +378,6 @@ def _dequantize(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                 )
                 # update the weight and remove the scale_inv tensor
                 state_dict[key] = dequantized_weight
-
-                # state_dict[key] = weight
                 scale_inv_keys.append(key + "_scale_inv")
 
         for key in scale_inv_keys:
@@ -411,10 +409,6 @@ def _add_quantization_scale_inv_tensors(
                 # add weight_scale_inv to the state_dict
                 weight_scale_inv_state_dict[key + "_scale_inv"] = torch.ones(
                     expected_scale_shape, dtype=torch.float32
-                )
-        print(
-            f"In _add_quantize_scale_inv_tensors, the added weight_scale_inv_state_dict: {weight_scale_inv_state_dict.keys()}"
-        )
 
         state_dict.update(weight_scale_inv_state_dict)
         return state_dict

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -47,13 +47,13 @@ dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 4
+tensor_parallel_degree = 1
 enable_async_tensor_parallel = false
 pipeline_parallel_degree = 1
 pipeline_parallel_schedule = "1F1B"
 context_parallel_degree = 1
-expert_parallel_degree = 2
-expert_tensor_parallel_degree = 4
+expert_parallel_degree = 1
+expert_tensor_parallel_degree = 1
 
 [checkpoint]
 enable = false

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_671b.toml

@@ -35,10 +35,10 @@ decay_type = "cosine"
 min_lr_factor = 0.1
 
 [training]
-local_batch_size = 2
-seq_len = 2048
+local_batch_size = 4
+seq_len = 4096
 max_norm = 1.0  # grad norm clipping
-steps = 10
+steps = 10_000
 compile = false
 dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
 
@@ -56,7 +56,7 @@ expert_tensor_parallel_degree = 1
 [checkpoint]
 enable = false
 folder = "checkpoint"
-interval = 10
+interval = 500
 last_save_model_only = true
 export_dtype = "float32"
 async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem]"