Dev/zero offload (#5858)

* fix llama * fix llama
5 months ago · 868afdb311
1 changed files with 0 additions and 2 deletions
--- a/colossalai/shardformer/modeling/llama.py
+++ b/colossalai/shardformer/modeling/llama.py
@ -803,8 +803,6 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
                shift_logits,
                shift_labels,
                process_group=shard_config.tensor_parallel_process_group,
-                vocab_size=self.lm_head.out_features,
-                dtype=self.model.dtype,
            )

        if not return_dict: