From 5db90c7f9b4f6a67a9b1c349dda372b15baeec17 Mon Sep 17 00:00:00 2001
From: Jiani Wang <jianiw@meta.com>
Date: Fri, 15 Aug 2025 17:35:44 -0700
Subject: [PATCH 1/2] free memory

---
 torchtitan/models/deepseek_v3/model/state_dict_adapter.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/torchtitan/models/deepseek_v3/model/state_dict_adapter.py b/torchtitan/models/deepseek_v3/model/state_dict_adapter.py
index 5a676b5a0..787bfe4f6 100644
--- a/torchtitan/models/deepseek_v3/model/state_dict_adapter.py
+++ b/torchtitan/models/deepseek_v3/model/state_dict_adapter.py
@@ -158,6 +158,9 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                     new_key = new_abstract_key.format(layer_num, expert_num)
                     hf_state_dict[new_key] = split_values[expert_num].squeeze()
 
+                # Remove the GroupedExperts' weight from the state_dict to free memory
+                del value
+
             elif "layers" in key:
                 abstract_key = re.sub(r"(\d+)", "{}", key, count=1)
                 layer_num = re.search(r"\d+", key).group(0)

From 672f73e1ed585f4e1e4af95e7948865c13299317 Mon Sep 17 00:00:00 2001
From: Jiani Wang <jianiw@meta.com>
Date: Sat, 16 Aug 2025 15:40:34 -0700
Subject: [PATCH 2/2] add README

---
 torchtitan/models/deepseek_v3/README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/torchtitan/models/deepseek_v3/README.md b/torchtitan/models/deepseek_v3/README.md
index 085403d47..92dd7bab7 100644
--- a/torchtitan/models/deepseek_v3/README.md
+++ b/torchtitan/models/deepseek_v3/README.md
@@ -61,6 +61,7 @@ python scripts/checkpoint_conversion/convert_from_hf.py <hf_checkpoints_dir> <dc
 Some limitations:
 1. It can't be used to convert HF checkpoint on the fly using GPU DTensor, because of sharding and quantized blocks may not be aligned well and causing silent numerfical incorrectness.
 2. It can't be used for weight sync to generate a state dict of bf16 because fake quantization to fp8 is applied.
+3. When converting GroupedExperts weights from HF separate expert weights on-the-fly, `torch.split()` will cause huge GPU memory usage. This is because torchtitan GroupedExperts' weight has shape `(num_experts, dim1, dim2)`, and by default shard FSDP on dim-0. When we call `torch.split()` in `to_hf()` function on dim-0, this will incur and all-gather and get replicated expert memory.
 
 ## To be added
 - Parallelism