Refactoring Trainer, adds `save_only_model` arg and simplifying FSDP integration #27652

pacman100 · 2023-11-22T12:02:11Z

What does this PR do?

Bumps up the minimum Accelerate version to 0.21.0
Add save_only_model arg - This enables the feature request Add an option to decide whether to store the checkpoint and rng_state. #26706
Simplifies a lot of logic in FSDP:
a. Currently, FSDP-XLA logic is custom in Trainer and normal FSDP is using the Accelerate's integration. There were many zombie code snippets related to normal FSDP. Cleaned those.
b. Made it easier to train with FSDP. When using FULL_STATE_DICT setting, it should now save the model in transformers format using the default safetensors sharded format. This reduces the burden on users to later load, shard and save in safetensors format.
c. Should fix Fine tuning with examples/pytorch/language-modeling/run_clm.py on torch/XLA + FSDP produce abnormal models #27432 but don't have access to TPUs to test this.
d. Fixes NotImplementedError: Cannot copy out of meta tensor; no data! #27166
e. This is built upon the PR in Accelerate to simplify FSDP integration fsdp refactoring accelerate#2177. It should be merged first.

1. Refactor FSDP 2. Add `--save_only_model` option: When checkpointing, whether to only save the model, or also the optimizer, scheduler & rng state. 3. Bump up the minimum `accelerate` version to `0.21.0`

HuggingFaceDocBuilderDev · 2023-11-22T12:22:42Z

The documentation is not available anymore as the PR was closed or merged.

This reverts commit 149330a.

muellerzr

Thanks a million! Looks great to me, just one tiny doc nit for consistency

src/transformers/training_args.py

Co-authored-by: Zach Mueller <muellerzr@gmail.com>

BenjaminBossan

Nice refactor and simplification of the code. I have a few comments, please take a look.

src/transformers/modeling_utils.py

BenjaminBossan · 2023-11-22T14:30:33Z

src/transformers/trainer.py

@@ -2462,12 +2396,49 @@ def _save_checkpoint(self, model, trial, metrics=None):
        else:
            torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth"))

-        if self.args.push_to_hub:
-            self._push_from_checkpoint(output_dir)
+    def _save_optimizer_and_scheduler(self, output_dir):


Hmm, it would be great if we could reuse accelerate code here, but it seems that there is no directly corresponding code. I wonder if we could achieve that with a bit of refactoring on the accelerate side @muellerzr (potentially in a future PR).

I looked into this the last time during migration to Accelerate backend as well and recently, I couldn't find a clean way of doing it on the Accelerate side. Maybe a discussion on this offline would help.

src/transformers/training_args.py

Co-authored-by: Benjamin Bossan <BenjaminBossan@users.noreply.github.com>

amyeroberts

Thanks for the refactor - looks a lot cleaner!

Regarding the accelerate bump up we should make sure everything's compatible in the library with this change. @ydshieh I believe the CI images should have the version reflected in the setup.py and modifying that file triggers a run on everything - is that right? Is there anything else we should check before merging?

amyeroberts · 2023-11-22T15:44:45Z

src/transformers/training_args.py

@@ -452,14 +459,18 @@ class TrainingArguments:
                    FSDP's limit_all_gathers (useful only when `fsdp` field is passed).
                     If `"True"`, FSDP explicitly synchronizes the CPU thread to prevent too many in-flight
                     all-gathers.
-                - use_orig_params (`bool`, *optional*, defaults to `False`)
+                - use_orig_params (`bool`, *optional*, defaults to `True`)


Why was this set to False by default before - is there an advantage to not having this enabled? Asking in case this introduces a degraded experience for some users

Hello, this is required now for simplifying the FSDP integration. Please find the explanation in the corresponding Accelerate PR: huggingface/accelerate#2177 (comment)

src/transformers/modeling_utils.py

src/transformers/trainer.py

amyeroberts

Sorry - meant to hit approve. Happy to marge after @BenjaminBossan's approval and @ydshieh confirms checking the library versions.

BenjaminBossan

Thanks for the updates. LGTM

ydshieh · 2023-11-23T13:13:28Z

Hi,

The (Circle)CI did run all the tests: for example, in torch_job, we can see

python -m pytest --junitxml=test-results/junit.xml -n 6 --max-worker-restart=0 --dist=loadfile --make-reports=tests_torch tests/benchmark tests/bettertransformer tests/deepspeed tests/extended tests/fixtures tests/fsdp tests/generation tests/models tests/optimization tests/peft_integration tests/quantization tests/sagemaker tests/test_backbone_common.py tests/test_configuration_common.py tests/test_configuration_utils.py tests/test_feature_extraction_common.py tests/test_feature_extraction_utils.py tests/test_image_processing_common.py tests/test_image_processing_utils.py tests/test_image_transforms.py tests/test_modeling_common.py tests/test_modeling_flax_common.py tests/test_modeling_flax_utils.py tests/test_modeling_tf_common.py tests/test_modeling_tf_utils.py tests/test_modeling_utils.py tests/test_pipeline_mixin.py tests/test_sequence_feature_extraction_common.py tests/test_tokenization_common.py tests/test_tokenization_utils.py tests/tokenization tests/tools tests/trainer tests/utils || true

The CI uses the latest main branch of accelerate (so newer than the latest released version)

"pip install -U --upgrade-strategy eager -e git+https://github.com/huggingface/accelerate@main#egg=accelerate"

Nothing extra to check.

welsh01 · 2023-11-28T17:56:57Z

save_only_model is a nice feature indeed, but it does not work together with load_best_model_at_end (at least with deepspeed enabled), since the final model cannot be loaded from the checkpoint.

pacman100 added 3 commits November 22, 2023 15:53

add code changes

eea69a6

1. Refactor FSDP 2. Add `--save_only_model` option: When checkpointing, whether to only save the model, or also the optimizer, scheduler & rng state. 3. Bump up the minimum `accelerate` version to `0.21.0`

quality

4d2a794

Merge branch 'main' into smangrul/FSDP-refactor

489a5df

pacman100 requested a review from muellerzr November 22, 2023 12:03

fix quality?

149330a

pacman100 added 3 commits November 22, 2023 18:01

Revert "fix quality?"

3dc64c9

This reverts commit 149330a.

fix fsdp doc strings

9685857

fix quality

203c38d

muellerzr approved these changes Nov 22, 2023

View reviewed changes

src/transformers/training_args.py Outdated Show resolved Hide resolved

pacman100 and others added 2 commits November 22, 2023 19:18

Update src/transformers/training_args.py

7e4f99a

Co-authored-by: Zach Mueller <muellerzr@gmail.com>

please fix the quality issue 😅

bebc321

pacman100 requested review from amyeroberts and BenjaminBossan November 22, 2023 13:49

BenjaminBossan requested changes Nov 22, 2023

View reviewed changes

pacman100 and others added 2 commits November 22, 2023 22:24

Apply suggestions from code review

835e41b

Co-authored-by: Benjamin Bossan <BenjaminBossan@users.noreply.github.com>

address comment

4237570

pacman100 mentioned this pull request Nov 22, 2023

remove the need to shard the model post saving when using FSDP pacman100/LLM-Workshop#13

Merged

amyeroberts reviewed Nov 22, 2023

View reviewed changes

amyeroberts approved these changes Nov 22, 2023

View reviewed changes

BenjaminBossan approved these changes Nov 23, 2023

View reviewed changes

pacman100 added 2 commits November 24, 2023 09:56

simplify conditional check as per the comment

bdef4ac

update documentation

c966e06

pacman100 merged commit a761d6e into main Nov 24, 2023
22 checks passed

pacman100 deleted the smangrul/FSDP-refactor branch November 24, 2023 06:10

pacman100 mentioned this pull request Dec 15, 2023

fix resuming from ckpt when using FSDP with FULL_STATE_DICT #27891

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactoring Trainer, adds `save_only_model` arg and simplifying FSDP integration #27652

Refactoring Trainer, adds `save_only_model` arg and simplifying FSDP integration #27652

pacman100 commented Nov 22, 2023

HuggingFaceDocBuilderDev commented Nov 22, 2023 •

edited

Loading

muellerzr left a comment

BenjaminBossan left a comment

BenjaminBossan Nov 22, 2023

pacman100 Nov 22, 2023

amyeroberts left a comment

amyeroberts Nov 22, 2023

pacman100 Nov 23, 2023

amyeroberts left a comment

BenjaminBossan left a comment

ydshieh commented Nov 23, 2023 •

edited

Loading

welsh01 commented Nov 28, 2023

Refactoring Trainer, adds save_only_model arg and simplifying FSDP integration #27652

Refactoring Trainer, adds save_only_model arg and simplifying FSDP integration #27652

Conversation

pacman100 commented Nov 22, 2023

What does this PR do?

HuggingFaceDocBuilderDev commented Nov 22, 2023 • edited Loading

muellerzr left a comment

Choose a reason for hiding this comment

BenjaminBossan left a comment

Choose a reason for hiding this comment

BenjaminBossan Nov 22, 2023

Choose a reason for hiding this comment

pacman100 Nov 22, 2023

Choose a reason for hiding this comment

amyeroberts left a comment

Choose a reason for hiding this comment

amyeroberts Nov 22, 2023

Choose a reason for hiding this comment

pacman100 Nov 23, 2023

Choose a reason for hiding this comment

amyeroberts left a comment

Choose a reason for hiding this comment

BenjaminBossan left a comment

Choose a reason for hiding this comment

ydshieh commented Nov 23, 2023 • edited Loading

welsh01 commented Nov 28, 2023

Refactoring Trainer, adds `save_only_model` arg and simplifying FSDP integration #27652

Refactoring Trainer, adds `save_only_model` arg and simplifying FSDP integration #27652

HuggingFaceDocBuilderDev commented Nov 22, 2023 •

edited

Loading

ydshieh commented Nov 23, 2023 •

edited

Loading