Tutorial for using DeepSpeed's activation checkpointing instead of PyTorch's #32409
Description
Feature request
Is there a tutorial for using DeepSpeed's activation checkpointing instead of PyTorch's?
I'm using Trainer with ZeRO integration to train my model. Here's my code:
if training_args.deepspeed_gradient_checkpointing and training_args.deepspeed: from deepspeed.runtime.activation_checkpointing.checkpointing import configure configure(mpu_=None) from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint model._set_gradient_checkpointing(training_args.deepspeed_gradient_checkpointing, checkpoint){ "activation_checkpointing": { "partition_activations": true, "cpu_checkpointing": true, "contiguous_memory_optimization": false, "number_checkpoints": null, "synchronize_checkpoint_boundary": false, "profile": false } }torchrun --nproc_per_node=8 \ --nnodes=${NNODES} \ --node_rank=${NODE_RANK} \ --master_addr=${MASTER_ADDR} \ --master_port=${MASTER_PORT} \ train.py \ --deepspeed ${DEEPSPEED_CONFIG_PATH} \ --gradient_checkpointing FalseHowever, I got this in FlashAttention2:
class XXXFlashAttention2(XXXAttention): def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: output_attentions = False bsz, q_len, _ = hidden_states.size() # <---- this got error query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) File "modeling_xxx.py", line 518, in forward bsz, q_len, _ = hidden_states.size() ValueError: not enough values to unpack (expected 3, got 2) Motivation
It seems there isn't such a tutorial available at the moment in either deepspeed's tutorial or huggingface.
Your contribution
Provide my results