feat(train): Support gradient checkpointing for Conformer & Transformer (whisper) #2173

xingchensong · 2023-11-27T12:56:25Z

2 brief intros on gradient checkpointing

TODO

benchmark
unit test

robin1001 · 2023-11-28T01:16:38Z

Reduce the memory in training, and the cost is slower training speed?

xingchensong · 2023-11-28T01:38:37Z

Reduce the memory in training, and the cost is slower training speed?

yes, "fit 10x larger neural nets into memory at the cost of an additional 20% computation time"

xingchensong · 2023-11-28T01:54:31Z

I think it's possible to fintune whisper-large-v3 (full parameter) on 24 * 2080ti, by Deepspeed + Grad Checkpointing + Flash Attention.

xingchensong · 2023-11-28T09:18:58Z

find_unused_parameters=False is required by activation checkpointing in DDP
find_unused_parameters=False is faster, according to
- https://pytorch-lightning.readthedocs.io/en/1.3.8/benchmarking/performance.html#when-using-ddp-set-find-unused-parameters-false
- Should Plugins expose certain parameters through Trainer? Lightning-AI/pytorch-lightning#6219
find_unused_parameters=True is required when ctc_weight == 0.0 or 1.0
- [fix] RuntimeError when ctc_weight=0.0/1.0 #206

Overall, we give a flag for users and defaulted it to False

xingchensong · 2023-11-29T05:20:57Z

benckmark on whisper-large-v3

4 * 3090 without gradient checkpointing, it will run into OOM

8 * 2080ti works .... amazing

config of deepspeed stage-3

{
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 4,
  "steps_per_print": 100,
  "gradient_clipping": 5,
  "fp16": {
    "enabled": false,
    "auto_cast": false,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "consecutive_hysteresis": false,
    "min_loss_scale": 1
  },
  "bf16": {
   "enabled": false
  },
  "zero_force_ds_cpu_optimizer": false,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "none",
      "pin_memory": true
    },
    "offload_param": {
      "device": "none",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 1e3,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 1e3,
    "contiguous_gradients" : true,
    "stage3_max_live_parameters": 2e3,
    "stage3_max_reuse_distance": 2e3,
    "stage3_prefetch_bucket_size": 1e3,
    "stage3_param_persistence_threshold": 1e3
  },
  "activation_checkpointing": {
    "partition_activations": false,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": null,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  },
  "flops_profiler": {
    "enabled": false,
    "profile_step": 100,
    "module_depth": -1,
    "top_modules": 1,
    "detailed": true,
    "output_file": null
  },
  "tensorboard": {
    "enabled": false,
    "output_path": "tensorboard/ds_logs/",
    "job_name": "deepspeed"
  }
}

xingchensong · 2023-11-29T10:10:12Z

benchmark on conformer

conformer.wenet.2.2.1: ddp
conformer.latest.373109c.grad_ckpt： ddp + grad_ckpt
conformer.latest.373109c.grad_ckpt.ds: deepspeed stage1 + grad_ckpt

cv_loss is almost identical to wenet v2.2.1
training with grad_ckpt has 15%~20% additional computation time compared to v2.2.1 (under the same train_config.yaml)

Mddct · 2023-11-29T13:12:09Z

fp32？

xingchensong · 2023-11-29T13:31:31Z

fp32？

yes ！！!

Mddct · 2023-11-29T13:33:02Z

amazing again！！！

xingchensong · 2023-11-30T06:07:35Z

TBD: gradient_checkpointing 应该在yaml中配置还是在train args中配置？目前是在yaml中配置的，不在args中配置，是考虑到目前只有transformer/conformer支持了grad_ckpt，其他模型都没支持，用户使用 squeezeformer 等模型时如果加了参数 --gradient_checkpoiting ，实际并不起作用（因为squeezeformer的forward里没实现这个）

xingchensong · 2023-11-30T08:22:50Z

pass unit test (gradient is equal), ready for a final review

xingchensong · 2023-11-30T09:06:42Z

wait, I just find that find_unused_param=False actually works in conjunction with ctc_weight=0.0. Perhaps we could simply eliminate the argument and retain find_unused_param=False as our default configuration.

Mddct · 2023-11-30T10:00:33Z

wait, I just find that find_unused_param=False actually works in conjunction with ctc_weight=0.0. Perhaps we could simply eliminate the argument and retain find_unused_param=False as our default configuration.

与模型有关的可能放yaml里会好些，也不想影响其他模块

btw：与train 有关的放arg里边

xingchensong · 2023-11-30T10:04:48Z

所以这里两个参数

find_unused_param ，加 or 不加？加的话加在哪合适
gradient_checkpointing，加 or 不加？加的话加在哪合适

Mddct · 2023-11-30T10:46:13Z

都放yanl里

xingchensong · 2023-11-30T12:13:22Z

known issue: can not pass JIT due to checkpoint function

  traceback : Traceback (most recent call last):
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
      return f(*args, **kwargs)
    File "/jfs-hdfs/user/xingchen.song/workspace/github/tmp/wenet.latest.grad_ckpt/wenet/examples/aishell/s0/wenet/bin/train.py", line 95, in main
      trace_and_print_model(args, model)
    File "/jfs-hdfs/user/xingchen.song/workspace/github/tmp/wenet.latest.grad_ckpt/wenet/wenet/utils/train_utils.py", line 372, in trace_and_print_model
      script_model = torch.jit.script(model)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_script.py", line 1286, in script
      return torch.jit._recursive.create_script_module(
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 476, in create_script_module
      return create_script_module_impl(nn_module, concrete_type, stubs_fn)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
      script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_script.py", line 615, in _construct
      init_fn(script_module)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 516, in init_fn
      scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
      create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 393, in create_methods_and_properties_from_stubs
      concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_recursive.py", line 863, in try_compile_fn
      return torch.jit.script(fn, _rcb=rcb)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/_script.py", line 1340, in script
      ast = get_jit_def(obj, obj.__name__)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/frontend.py", line 293, in get_jit_def
      return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/frontend.py", line 331, in build_def
      param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/jit/frontend.py", line 355, in build_param_list
      raise NotSupportedError(ctx_range, _vararg_kwarg_err)
  torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
    File "/bucket/output/jfs-hdfs/user/xingchen.song/tools/miniconda3/lib/python3.9/site-packages/torch/utils/checkpoint.py", line 164
  def checkpoint(function, *args, use_reentrant: bool = True, **kwargs):
                                                               ~~~~~~~ <--- HERE
      r"""Checkpoint a model or part of the model

BenBill2077 · 2023-11-30T17:26:44Z

What is gradient ckpt？

xingchensong · 2023-12-01T01:49:47Z

What is gradient ckpt？

一句话总结: 2080ti 全量参数微调whisper-large 1.5B 不再是梦😜（具体技术细节可见开头的blog 链接）

xingchensong · 2023-12-01T06:32:05Z

pass JIT ! please review this PR, @robin1001 @Mddct

xingchensong mentioned this pull request Nov 27, 2023

中文开源语音大模型计划 #2097

Open

14 tasks

xingchensong force-pushed the xcsong-grad-ckpt branch from 3b318ec to fbb0557 Compare November 28, 2023 07:28

xingchensong changed the title ~~feat(train): Support gradient checkpointing for Conformer & Transformer~~ feat(train): Support gradient checkpointing for Conformer & Transformer (whisper) Nov 28, 2023

xingchensong force-pushed the xcsong-grad-ckpt branch 3 times, most recently from 79e8f2f to 0bbff4a Compare November 29, 2023 03:55

Mddct self-requested a review November 29, 2023 05:38

xingchensong mentioned this pull request Nov 29, 2023

refactor(deepspeed): Refine traning code #2055

Merged

6 tasks

xingchensong requested review from robin1001 and whiteshirt0429 November 30, 2023 08:23

xingchensong added 2 commits December 1, 2023 14:01

feat(train): Support gradient checkpointing for Conformer & Transformer

c2265a7

feat(train): config grad-ckpt in yaml

b424fe3

xingchensong added 6 commits December 1, 2023 14:01

feat(train): training success

43099b1

feat(train): add note

a644cf2

feat(train): set grad_ckpt=True for whisper

aa3fc93

feat(train): add unit test for grad_ckpt

a6ab4d9

feat(train): remove find_unused_parameter

3786e4a

feat(train): pass jit.script

0ecbd37

xingchensong force-pushed the xcsong-grad-ckpt branch from 2be317c to 0ecbd37 Compare December 1, 2023 06:01

xingchensong added 3 commits December 1, 2023 14:06

feat(train): add unit test for jit.script

2a6df3b

feat(train): add unit test for jit.script

4d13a24

feat(train): fix unit test for grad_ckpt

3f1d2f0

Mddct approved these changes Dec 1, 2023

View reviewed changes

Mddct merged commit 8d6d23f into main Dec 1, 2023
6 checks passed

Mddct deleted the xcsong-grad-ckpt branch December 1, 2023 07:25

xingchensong mentioned this pull request Dec 8, 2023

[fix] set find_unused_parameters=False only when enable grad_ckpt #2215

Merged

xingchensong mentioned this pull request Dec 20, 2023

[examples] whisper aishell #2238

Merged

Mddct mentioned this pull request Mar 27, 2024

[train_engine] support fsdp #2412

Merged

24 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(train): Support gradient checkpointing for Conformer & Transformer (whisper) #2173

feat(train): Support gradient checkpointing for Conformer & Transformer (whisper) #2173

xingchensong commented Nov 27, 2023 •

edited

Loading

robin1001 commented Nov 28, 2023

xingchensong commented Nov 28, 2023

xingchensong commented Nov 28, 2023

xingchensong commented Nov 28, 2023

xingchensong commented Nov 29, 2023 •

edited

Loading

xingchensong commented Nov 29, 2023 •

edited

Loading

Mddct commented Nov 29, 2023

xingchensong commented Nov 29, 2023

Mddct commented Nov 29, 2023

xingchensong commented Nov 30, 2023

xingchensong commented Nov 30, 2023

xingchensong commented Nov 30, 2023 •

edited

Loading

Mddct commented Nov 30, 2023 •

edited

Loading

xingchensong commented Nov 30, 2023 •

edited

Loading

Mddct commented Nov 30, 2023

xingchensong commented Nov 30, 2023 •

edited

Loading

BenBill2077 commented Nov 30, 2023

xingchensong commented Dec 1, 2023 •

edited

Loading

xingchensong commented Dec 1, 2023

feat(train): Support gradient checkpointing for Conformer & Transformer (whisper) #2173

feat(train): Support gradient checkpointing for Conformer & Transformer (whisper) #2173

Conversation

xingchensong commented Nov 27, 2023 • edited Loading

TODO

robin1001 commented Nov 28, 2023

xingchensong commented Nov 28, 2023

xingchensong commented Nov 28, 2023

xingchensong commented Nov 28, 2023

xingchensong commented Nov 29, 2023 • edited Loading

benckmark on whisper-large-v3

config of deepspeed stage-3

xingchensong commented Nov 29, 2023 • edited Loading

benchmark on conformer

Mddct commented Nov 29, 2023

xingchensong commented Nov 29, 2023

Mddct commented Nov 29, 2023

xingchensong commented Nov 30, 2023

xingchensong commented Nov 30, 2023

xingchensong commented Nov 30, 2023 • edited Loading

Mddct commented Nov 30, 2023 • edited Loading

xingchensong commented Nov 30, 2023 • edited Loading

Mddct commented Nov 30, 2023

xingchensong commented Nov 30, 2023 • edited Loading

BenBill2077 commented Nov 30, 2023

xingchensong commented Dec 1, 2023 • edited Loading

xingchensong commented Dec 1, 2023

xingchensong commented Nov 27, 2023 •

edited

Loading

xingchensong commented Nov 29, 2023 •

edited

Loading

xingchensong commented Nov 29, 2023 •

edited

Loading

xingchensong commented Nov 30, 2023 •

edited

Loading

Mddct commented Nov 30, 2023 •

edited

Loading

xingchensong commented Nov 30, 2023 •

edited

Loading

xingchensong commented Nov 30, 2023 •

edited

Loading

xingchensong commented Dec 1, 2023 •

edited

Loading