Fine-tuning T5 with long sequence length, using activation checkpointing with Deepspeed
I’m fine-tuning T5 (11B) with very long sequence lengths (2048 input, 256 output) and am running out of memory on an 8x A100-80GB cluster even with ZeRO-3 enabled, bf16 enabled, and per-device batch size=1. The issue seems to be not with optimizer or model memory, but rather activation memory. I’m trying to get activation checkpointing to work with my existing setup (which uses the automatic HF Trainer/Deepspeed integration).Would really appreciate your advice@stas
2022-12-05T02:51:20Z
Indeed, enabling activation checkpointing should make a very noticeable difference.If that is not enough you can look intoMemory-centric tilingwhich should shave some more memory, and tuning up buffer sizes in the deepspeed config may help a bit more.Specifically to your situationSequence Parallelismshould be very helpful, but if I’m not mistaken this is yet to be supported by Deepspeed.- you may want to submit a feature request for this to happen.The frameworks that support SP that I know of areMegatron-LM,CollosalAI,Transformer Engine. There might be others.
2022-12-05T03:24:54Z
Thank you,@stas.My issue is that enabling activation checkpointing in the ds_config makes no difference in terms of used memory (as measured by nvidia-smi).Specifically, I’ve added"activation_checkpointing": {
    "partition_activations": True,
    }But it seems to be useless. Would appreciate your input.
2022-12-05T03:41:36Z
Ah, yes, that feature is part of the modeling code. Please see:huggingface.coPerformance and Scalability: How To Fit a Bigger Model and Train It FasterWe’re on a journey to advance and democratize artificial intelligence through open source and open science.the example is for the HF Trainer, but you can, of course, do the same w/o the Trainer. seegithub.comhuggingface/transformers/blob/699e90437f984d69ad3c9b891dd2e9d0fc2cffe4/src/transformers/modeling_utils.py#L1465to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 onlayer 1 and heads 2 and 3 on layer 2."""# save new sets of pruned heads as union of previously stored pruned heads and newly pruned headsfor layer, heads in heads_to_prune.items():union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads)self.config.pruned_heads[layer] = list(union_heads)  # Unfortunately we have to store it as list for JSONself.base_model._prune_heads(heads_to_prune)def gradient_checkpointing_enable(self):"""Activates gradient checkpointing for the current model.Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpointactivations"."""if not self.supports_gradient_checkpointing:raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")self.apply(partial(self._set_gradient_checkpointing, value=True))it’s confusing that there are 2 very different names for the same feature in the ml world.
2022-12-05T04:43:38Z
Ah I see, so activation and gradient checkpointing are the same thing? The Deepspeed activation checkpoint referenceseems to suggestthat their implementation partitions the activations between the GPUs (similar to gradients + model weights in ZeRO 3).Does this gradient_checkpointing=True flag HF Trainer enable partitioning as well? That is an optimization I’m interested in – as most of my GPU memory is in fact being eaten up by activations.
2022-12-05T06:23:48Z
Additionally, when setting gradient_checkpointing=True with distributed multi-node Deepspeed (4 x 8xA100), I get constant warnings@stas:use_cache=Trueis incompatible with gradient checkpointing. Settinguse_cache=False…Thanks for your help!
2022-12-05T17:10:35Z
It tells you that it sets it toFalse, if you don’t want the warming, set theuse_cache=Falseexplicitlygithub.comhuggingface/transformers/blob/28d74872cc049e0cbee3fafd15cbbabfe348ebd4/src/transformers/generation_utils.py#L689-L691use_cache: (:obj:`bool`, `optional`, defaults to :obj:`True`):Whether or not the model should use the past last key/values attentions (if applicable to the model) tospeed up decoding.Incidentally, one can turn the caching off duringgenerateto save memory, but at a cost of regenerating the past values.