fixed 0 layers

argonne-lcf · Feb 6, 2025 · d25d730 · d25d730
1 parent 02b31e9
commit d25d730
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dlio_benchmark/checkpointing/base_checkpointing.py b/dlio_benchmark/checkpointing/base_checkpointing.py
@@ -168,7 +168,7 @@ def get_layer_parameters(self, layer_index):
                 sharding_factor = self.dp
             h, l, ffn, voc = self.args.hidden_size, self.args.num_layers, self.args.ffn_hidden_size, self.args.vocab_size
             if layer_index == 0 or layer_index == l + 1:
-                return [h * voc // sharding_factor * dtype_size] # embedding or lm_head
+                return [h * voc // self.tp // sharding_factor * dtype_size] # embedding or lm_head
             elif layer_index == l + 2:
                 return [h//sharding_factor * dtype_size]
             else: