Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: add unified reshard script #734

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add unified reshard script #734

wants to merge 1 commit into from

Conversation

mattmazzola
Copy link
Contributor

@mattmazzola mattmazzola commented Jun 9, 2023 •
edited

Issue

Training requires flattened models, any MP and FSDP
Inference requires unflattened models with FSDP 1

We wanted AML jobs which train model (which produces Flattened checkpoint output), reshards the checkpoints to Unflatten with FSDP 1 for inference, and evaluates the model outputs

  • Did not have a way to convert flattened checkpoints from one model parallel version to another
  • Did not have a way to convert an inference checkpoint to a training checkpoint
  • Converting a trained checkpoint for inference requires several steps

Solution

  • Add unified reshard script
    • Allows modifying both MP and FSDP in single script
    • Supports flattening/unflattening the model files
    • Automatically infers and adds shard_metadata
    • Improved logging output to verify expect transformation

Testing steps

Did not test although this file was direct from our fork without modification so it likely works as is

Video or Screenshots

Flow Chart of Steps in Script

MPxFSDP (2x8 -> 8x1)

2023-06-01 19:06:53 | INFO | __main__ | Found 16 sharded checkpoints (checkpoint_last-model_part-0-shard0.pt to checkpoint_last-model_part-1-shard7.pt)
2023-06-01 19:06:53 | INFO | __main__ | Loading all sharded checkpoints to CPU
2023-06-01 19:07:01 | INFO | __main__ | Detected Input Model State:
2023-06-01 19:07:01 | INFO | __main__ | - Model Parallel (MP) factor:			2
2023-06-01 19:07:01 | INFO | __main__ | - Fully Sharded Data Parallel (FSDP) factor:	8
2023-06-01 19:07:01 | INFO | __main__ | - Model Weights:				Flattened
2023-06-01 19:07:01 | INFO | __main__ | Desired Output Model State:
2023-06-01 19:07:01 | INFO | __main__ | - Model Parallel (MP) factor:			8
2023-06-01 19:07:01 | INFO | __main__ | - Fully Sharded Data Parallel (FSDP) factor:	1
2023-06-01 19:07:01 | INFO | __main__ | - Model Weights:				Unflattened
2023-06-01 19:07:01 | INFO | __main__ | Resharding model parallel from 2 to 8
2023-06-01 19:07:01 | INFO | __main__ | You attempted to change MP from 2 to 8 but the models was flattened. It must be unflattened with FSDP 1 to change model parallel factor. Unflattening and consolidating the weights.
2023-06-01 19:07:01 | INFO | __main__ | Resharding state dicts into 1 fsdp shard(s)
2023-06-01 19:07:01 | INFO | __main__ | Resharding model weights into 1 shard(s)
2023-06-01 19:07:02 | INFO | __main__ | Resharding state dicts into 1 fsdp shard(s)
2023-06-01 19:07:02 | INFO | __main__ | Resharding model weights into 1 shard(s)
2023-06-01 19:07:02 | INFO | __main__ | Allocating memory for unsharded checkpoint
2023-06-01 19:07:03 | WARNING | __main__ | Max value discrepancy for key 'decoder.embed_positions.weight': 4.7684e-06
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 0
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 1
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 2
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 3
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 4
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 5
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 6
2023-06-01 19:07:04 | INFO | __main__ | Resharding state dict for model parallel part 7
2023-06-01 19:07:05 | INFO | __main__ | n_layers: 24
2023-06-01 19:07:06 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-0.pt
2023-06-01 19:07:10 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-1.pt
2023-06-01 19:07:14 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-2.pt
2023-06-01 19:07:18 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-3.pt
2023-06-01 19:07:23 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-4.pt
2023-06-01 19:07:27 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-5.pt
2023-06-01 19:07:31 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-6.pt
2023-06-01 19:07:35 | INFO | __main__ | Writing a resharded state dict to //amltbf97debed90d3226669ab66b89d50d4a/projects/mattm-nlg-distill/amlt-results/7314329603.77649-0eb9a9d5-5ed5-4908-9e39-7cafc2438bcd/opt-1.3b-e2e_nlg-human--full-20230601.1144/reshard/reshard-model_part-7.pt

Flattened to Unflattened Output (MP and FSDP unchanged):

Found 2 sharded checkpoints (checkpoint_last-model_part-0-shard0.pt to checkpoint_last-model_part-1-shard0.pt)
Loading all sharded checkpoints to CPU
Detected Input Model State:
- Model Parallel (MP) factor:                   2
- Fully Sharded Data Parallel (FSDP) factor:    1
- Model Weights:                                Flattened
Desired Output Model State:
- Model Parallel (MP) factor:                   2
- Fully Sharded Data Parallel (FSDP) factor:    1
- Model Weights:                                Unflattened
Current model Flattened but desired model is Unflattened
Unflattened model
Resharding state dicts into 1 fsdp shard(s)
Resharding model weights into 1 shard(s)
Resharding state dicts into 1 fsdp shard(s)
Resharding model weights into 1 shard(s)
Writing a resharded state dict to _results/1.3b-2x1-unflattened/reshard-model_part-0-shard0.pt
Writing a resharded state dict to _results/1.3b-2x1-unflattened/reshard-model_part-1-shard0.pt

Related to #726

Most of work was done by @sahajgg

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@suchenzang suchenzang Awaiting requested review from suchenzang

@ngoyal2707 ngoyal2707 Awaiting requested review from ngoyal2707 ngoyal2707 is a code owner

@punitkoura punitkoura Awaiting requested review from punitkoura punitkoura is a code owner

@moyapchen moyapchen Awaiting requested review from moyapchen

@klshuster klshuster Awaiting requested review from klshuster

@davides davides Awaiting requested review from davides davides is a code owner

@igormolybogFB igormolybogFB Awaiting requested review from igormolybogFB

@Xirider Xirider Awaiting requested review from Xirider

@andrewPoulton andrewPoulton Awaiting requested review from andrewPoulton

@bashnick bashnick Awaiting requested review from bashnick bashnick is a code owner

@tangbinh tangbinh Awaiting requested review from tangbinh tangbinh is a code owner

@urielsinger urielsinger Awaiting requested review from urielsinger

@zycalice zycalice Awaiting requested review from zycalice zycalice is a code owner

@ArmenAg ArmenAg Awaiting requested review from ArmenAg ArmenAg is a code owner

@lilisierrayu lilisierrayu Awaiting requested review from lilisierrayu lilisierrayu is a code owner

@adampolyak adampolyak Awaiting requested review from adampolyak adampolyak is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
cla signed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
2 participants
@mattmazzola @facebook-github-bot