fix mixed precision for `replicate` / pure DDP #591

152334H · 2024-09-29T07:03:42Z

Hi. I noticed the following:

the keyword autocast does not exist in the repository
MixedPrecisionConfig is only used in the fully_shard codepath
the duration of a dummy 1000 step run is a lot longer with DDP than with FSDP

All of the above indicates that, when dp_shard_enabled is false, training runs with pure fp32, regardless of the mixed precision config.

This pull request changes the code to use torch.autocast in the training forward pass, specifically only when dp_shard_enabled is false, to the dtype of mixed_precision_param.

152334H · 2024-09-29T07:28:23Z

If mixed_precision_reduce ever has arguments beyond float32, it may be worth adding an appropriate compression hook after replicate as well, e.g.

diff --git a/torchtitan/parallelisms/parallelize_llama.py b/torchtitan/parallelisms/parallelize_llama.py
index fc26703..bdbcbb2 100644
--- a/torchtitan/parallelisms/parallelize_llama.py
+++ b/torchtitan/parallelisms/parallelize_llama.py
@@ -344,5 +344,7 @@ def apply_ddp(
             torch._dynamo.config.optimize_ddp = "ddp_optimizer"

     replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100)
+    from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import bf16_compress_hook
+    model.register_comm_hook(dp_mesh.get_group(), bf16_compress_hook)

     logger.info("Applied DDP to the model")

awgu · 2024-09-29T15:49:32Z

might be nice to fold this as part of the train_context

152334H · 2024-09-29T18:32:09Z

@awgu it can't be folded into the train context unless you're certain autocasting the backward will lead to no problems. According to torch documentation,

autocast should wrap only the forward pass(es) of your network, including the loss computation(s). Backward passes under autocast are not recommended. Backward ops run in the same type that autocast used for corresponding forward ops.

Personally, in other projects, I have experienced issues where torch.compiling a network with autocast over the backwards inexplicably causes NaNs to appear. I don't know whether that applies in this case but will take any overruling if people are certain it will work without problems.

awgu · 2024-09-30T12:00:51Z

cc: @tianyu-l on thoughts on how to handle this
perhaps separate forward and backward contexts

tianyu-l · 2024-09-30T19:43:50Z

Hmm I didn't have much context on this.

@fegin How is DDP supposed to handle mixed precision? Is AMP autocasting suggested?

awgu · 2024-10-01T03:11:09Z

I think Rohan added a tentative mixed precision API for DDP, but it never made it to public feature. I think using AMP is probably the way to go.

fegin · 2024-10-01T17:51:58Z

autocast is the right way for DDP.

tianyu-l · 2024-10-04T00:07:16Z

torchtitan/config_manager.py

-                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`.
+                When data_parallel_degree == 1, this enables AMP autocast.


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

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.

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

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

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

DTensor TP+ AMP

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

tianyu-l · 2024-10-04T00:23:22Z

train.py

+                    with contextlib.nullcontext() if parallel_dims.dp_shard_enabled else torch.autocast(
+                        "cuda",
+                        dtype=TORCH_DTYPE_MAP[
+                            job_config.training.mixed_precision_param
+                        ],
+                    ):


IMO there are two ways:

To simplify it here, we can have a separate get_foward_only_context() which encapsulates the logic.

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?

2 sounds good to me

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

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.

torchtitan/config_manager.py

152334H added 2 commits September 29, 2024 14:58

use autocast for replicate

62d0d95

update mixed_precision_param desc

dae0527

facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 29, 2024

awgu requested a review from tianyu-l September 30, 2024 12:00

possibly fix lint error

4da7376

tianyu-l requested a review from fegin October 3, 2024 18:23

tianyu-l reviewed Oct 4, 2024

View reviewed changes

Update config_manager.py

747c3b5

tianyu-l mentioned this pull request Oct 18, 2024

Non-DP runs default to float32 precision #630

Open

tianyu-l linked an issue Nov 22, 2024 that may be closed by this pull request

Non-DP runs default to float32 precision #630

Open

garrett361 mentioned this pull request Nov 25, 2024

Is autocast needed with FSDP2? #700

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix mixed precision for `replicate` / pure DDP #591

fix mixed precision for `replicate` / pure DDP #591

152334H commented Sep 29, 2024

152334H commented Sep 29, 2024

awgu commented Sep 29, 2024 •

edited

Loading

152334H commented Sep 29, 2024

awgu commented Sep 30, 2024

tianyu-l commented Sep 30, 2024

awgu commented Oct 1, 2024

fegin commented Oct 1, 2024

tianyu-l Oct 4, 2024

152334H Oct 4, 2024

tianyu-l Oct 4, 2024

fegin Oct 8, 2024

152334H Oct 8, 2024

fegin Oct 8, 2024

tianyu-l Oct 18, 2024

tianyu-l Oct 4, 2024

awgu Oct 4, 2024

152334H Oct 4, 2024

tianyu-l Oct 4, 2024

fegin Oct 8, 2024

fix mixed precision for replicate / pure DDP #591

Are you sure you want to change the base?

fix mixed precision for replicate / pure DDP #591

Conversation

152334H commented Sep 29, 2024

152334H commented Sep 29, 2024

awgu commented Sep 29, 2024 • edited Loading

152334H commented Sep 29, 2024

awgu commented Sep 30, 2024

tianyu-l commented Sep 30, 2024

awgu commented Oct 1, 2024

fegin commented Oct 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fix mixed precision for `replicate` / pure DDP #591

fix mixed precision for `replicate` / pure DDP #591

awgu commented Sep 29, 2024 •

edited

Loading