Skip to content

Commit f931aa9

Browse files
committed
enable DSV3 manual bucketing
1 parent 605a9a1 commit f931aa9

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,11 @@ def convert_modules_to_fqns(modules, module_to_fqn_mapping):
3838
result = []
3939
for m in modules:
4040
if isinstance(m, list):
41-
result.append(convert_modules_to_fqns(m, module_to_fqn_mapping))
41+
if fqn_list := convert_modules_to_fqns(m, module_to_fqn_mapping):
42+
result.append(fqn_list)
4243
else:
43-
result.append(module_to_fqn_mapping.get(m, None))
44+
if fqn := module_to_fqn_mapping.get(m):
45+
result.append(fqn)
4446
return result
4547

4648
module_to_name = {m: n for n, m in model.named_modules()}

0 commit comments

Comments
 (0)