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

fix mixed precision for replicate / pure DDP #591

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix mixed precision for replicate / pure DDP #591

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
62d0d95
use autocast for replicate
152334H Sep 29, 2024
dae0527
update mixed_precision_param desc
152334H Sep 29, 2024
4da7376
possibly fix lint error
152334H Sep 30, 2024
747c3b5
Update config_manager.py
152334H Oct 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions torchtitan/config_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,8 +332,9 @@ def __init__(self):
default="bfloat16",
choices=["bfloat16", "float32"],
help="""
torch dtype to use for parameters when applying mixed precision via FSDP.
This feature only takes effect when data_parallel_degree > 1
torch dtype to use for parameters when applying mixed precision.
When data_parallel_degree > 1, this changes FSDP's `param_dtype`.
152334H marked this conversation as resolved.
Show resolved Hide resolved
When data_parallel_degree == 1, this enables AMP autocast.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit confused by the description here. Other than DDP, does it work with TP, PP, CP, etc.?
I wonder if there is a document/tutorial on this.
cc: @fegin

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the implementation of mixed precision only needs to change based on whether fully_shard (which has its own internal mixed prec mechanism) or replicate (which internally wraps a module with DDP) is used.

therefore, the qn of what happens (and what should happen) w/ TP/PP/DP is solely dependent on the sharded data parallel degree.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what happens if neither FSDP or DDP is used? btw DDP is not composable with TP/PP today as far as I know.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If FSDP is not used, AMP should be the default answer for all other parallelisms. But we don't actually test AMP + TP.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you mean the behavior of dtensor TP sharded operations under AMP in general pytorch is not tested? or that specifically it is not tested in torchtitan

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DTensor TP+ AMP

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@152334H Would you be able to test TP + AMP to see if it works? Maybe adding a unit test would be helpful.

""",
)
self.parser.add_argument(
Expand Down
12 changes: 9 additions & 3 deletions train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

from torchtitan import utils
from torchtitan.checkpoint import CheckpointManager, TrainState
from torchtitan.config_manager import JobConfig
from torchtitan.config_manager import JobConfig, TORCH_DTYPE_MAP
from torchtitan.datasets import build_hf_data_loader, build_tokenizer
from torchtitan.float8 import Float8Handler
from torchtitan.logging import init_logger, logger
Expand Down Expand Up @@ -288,8 +288,14 @@ def loss_fn(pred, labels):
else:
# Non-PP forward / backward
with train_context():
pred = model(input_ids)
loss = loss_fn(pred, labels)
with contextlib.nullcontext() if parallel_dims.dp_shard_enabled else torch.autocast(
"cuda",
dtype=TORCH_DTYPE_MAP[
job_config.training.mixed_precision_param
],
):
Comment on lines +291 to +296
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMO there are two ways:

  1. To simplify it here, we can have a separate get_foward_only_context() which encapsulates the logic.
  2. Modify get_train_context to receive a boolean variable is_backward so that we can obtain two different contexts, one for forward, the other for backward. Pro: only one method to obtain context. Con: need to enter and exit twice for context managers shared by forward and backward.
    cc: @awgu wdyt?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 sounds good to me

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will maybe_enable_compiled_autograd work correctly when it exits and re-enters in-between a fwd-bck?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@fegin

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ye, that should work. It's a context manager to enable the flag. So it should be okay to separately wrap fwd and bwd with maybe_enable_compiled_autograd.

pred = model(input_ids)
loss = loss_fn(pred, labels)
# pred.shape=(bs, seq_len, vocab_size)
# need to free to before bwd to avoid peaking memory
del pred
Expand Down
Loading