Skip to content

GitLab

Switch branch/tag
  1. 05 Mar, 2023 3 commits
    • Fei Sun's avatar
      Move EMA to after backward. · a7dc757c
      Fei Sun authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/494

Currently EMA computation is in the after step hook. It is in the critical path where no other work is available. This increases the training iteration time. This diff moves the EMA computation to after the backward but before the optimizer step. This way, the majority of the EMA computation time on the CPU can be hidden since CPU at that time is waiting for the GPU to finish the backward anyway. This change may completely hide the EMA CPU time. It reduces the EMA time from 20ms to 4ms, where the 4ms is the GPU time.

However, with this change, the EMA gets its value from the previous iteration value (since it is before step). but since we do many epochs of training, one iteration difference may not be significant.

Reviewed By: tglik

Differential Revision: D43527552

fbshipit-source-id: 1faa9d910b20cae0fc77da541bc0ad176bce18a8
      a7dc757c
    • Fei Sun's avatar
      Prefetch forward · 5f1ef548
      Fei Sun authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/492

Enable prefetching the FSDP all gathers. Forward prefetch may or may not improve performance. Its effectiveness is determined by other FSDP options, such as zero2/zero3, HSDP/FSDP. Need to do a HPO sweep to figure out the best configuration.

Reviewed By: wat3rBro

Differential Revision: D43027253

fbshipit-source-id: cbf1b4bcf5b0b8301b5b9547e3c22b1f0ffc7590
      5f1ef548
    • Fei Sun's avatar
      Use LERP to implement EMA · 255313d8
      Fei Sun authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/493

Currently the EMA implementation first does the multiplication and then does the addition. It requires two round trips from HBM. With the lerp operator, one kernel can do both. This change uses LERP to compute EMA instead. It reduces the GPU EMA computation time by 40%.

Reviewed By: newstzpz

Differential Revision: D43525938

fbshipit-source-id: ca1e14453bdfda958d3c412a52ff48efa65b3dd4
      255313d8
  3. 02 Mar, 2023 1 commit
    • Anthony Chen's avatar
      add correctness test for ai infra checkpointer with FSDP local mode in d2go runner · fd0cbb8f
      Anthony Chen authored 
      Summary:
X-link: https://github.com/facebookresearch/mobile-vision/pull/141

Pull Request resolved: https://github.com/facebookresearch/d2go/pull/484

This diff adds a correctness test for ai infra checkpointer + FSDP local mode within a d2go runner. It verifies that ai infra checkpointer saves the exact same model as the old checkpointer, and that we can convert between ai infra checkpoints (local) and fsdp checkpoints (local + global) seamlessly.

Note: adapted from mattcyu1's script D43492498.

## Testing

Testing is done by saving with both ai infra and fsdp checkponter and compare the state dict produced. Here are the steps:

1. Build the model. Save a local ckpt using the FSDP checkpointer and another local ckpt using the AIInfra checkpointer
2. Reset the model. Load local ckpt using the FSDP checkpointer and convert it to global ckpt
3. Reset the model. Load local ckpt using the AIInfra checkpointer and re-save it as global ckpt using the FSDP checkpointer
4. Compare the two global state dicts

## Others

1. Add a launch decorator for d2go.distributed worker using the one in `fbcode/mobile-vision/mobile_cv/mobile_cv/torch/utils_pytorch/distributed_helper.py`

2. Remove `ema_state.load_state_dict()` in loading. This is needed because ai infra checkpointer loads state dict in place before `ema_state.load_state_dict()` is called. Since it's inplace loading, both ema_state and state_dict['ema_state'] points to the same tensor. Calling` ema.load_state_dict()` clears ema_state, effectively freeing the tensor and cause it to return an empty dict.
Solution: Don't call `ema_state.load_state_dict()` because it's already loaded. More info: https://www.internalfb.com/intern/wiki/Components_in_AI/Checkpoint/Getting_Started/Input_Output_Contract/#load

Reviewed By: xunnanxu

Differential Revision: D43423572

fbshipit-source-id: 8c4a47917670ea1205f952540d1e4cb9fc9232c0
      fd0cbb8f
  5. 01 Mar, 2023 2 commits
  7. 28 Feb, 2023 1 commit
  9. 25 Feb, 2023 1 commit
  11. 24 Feb, 2023 1 commit
    • Matthew Yu's avatar
      turn off interleaving if only saving on rank0 · 3111ae59
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/482

We should avoid using interleaving during save if we are calling save on one process:
```
if comm.is_main_process():
  save()
```
this is because interleave calls comm.synchronize() so will just wait indefinitely.

This diff updates the FSDP checkpointer to use save(interleave=False) when running on one process.

Reviewed By: wat3rBro, YanjunChen329

Differential Revision: D43526328

fbshipit-source-id: 672993a87af627aca090384b0c218798bd42fcde
      3111ae59
  13. 23 Feb, 2023 2 commits
  15. 17 Feb, 2023 2 commits
  17. 16 Feb, 2023 4 commits
    • Anthony Chen's avatar
      Add an option to specify the period of metric gathering and writing in Trainer · 6f43a43a
      Anthony Chen authored 
      Summary:
X-link: https://github.com/fairinternal/detectron2/pull/591

Pull Request resolved: https://github.com/facebookresearch/d2go/pull/469

X-link: https://github.com/facebookresearch/detectron2/pull/4785

Add an option to specify the period of metric gathering and writing in Trainer.

This feature is needed to optimize training speed for large-scale training jobs like generative AI. The reason is that the all_gather call in metric writing at every iteration is time-consuming when hundreds of gpus are used. This takes ~10% of the total training time. With this feature we can set the metric writing period as the same as cfg.WRITER_PERIOD=20 to reduce training time while still keeping metric logging the same to users

Reviewed By: miqueljubert, wat3rBro

Differential Revision:
D43098985

Privacy Context Container: 2011691122555468

fbshipit-source-id: 63c93a7331aa63badce5125e5240d2d5f7e61b74
      6f43a43a
    • Sudarshan Raghunathan's avatar
      Add reply files to d2go training processes · f0f55cdc
      Sudarshan Raghunathan authored 
      Summary:
This diff contains a minimal set of changes to support returning reply files to MAST.

There are three parts:
1. First, we have a try..except in the main function to catch all the "catchable" Python exceptions. Exceptions from C++ code or segfaults will not be handled here.
2. Each exception is then written to a per-process JSON reply file.
3. At the end, all per-process files are stat-ed and the earliest file is copied to a location specified by MAST.

# Limitations
1. This only works when local processes are launched using multiprocessing (which is the default)
2. If any error happens in C++ code - it will likely not be caught in Python and the reply file might not have the correct logs

Differential Revision: D43097683

fbshipit-source-id: 0eaf4f19f6199a9c77f2ce4c7d2bbc2a2078be99
      f0f55cdc
    • Tao Xu's avatar
      fix the issue of tensorboard visualization · b21607b1
      Tao Xu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/473

As shown in the attached image and tb visualization, some of our jobs fail to save the results to tensorboard.
There should be some messages between circled lines of the screenshot if the images are added to tensorboard.
One possible reason is that the tensorbord visualization evaluator is only added for the rank 0 gpu. It may fail to fetch any data during evaluation of diffusion model which only do 1 batch of inference during validataion.
To resolve this issue, we add the visualization evaluator to all ranks of gpus and gather their results, and finally add the results with biggest batchsize to the tensorboard for visualization.

The screenshot is from f410204704 (https://www.internalfb.com/manifold/explorer/mobile_vision_workflows/tree/workflows/xutao/20230211/latest_train/dalle2_decoder.SIULDLpgix/e2e_train/log.txt)

Refactored the default_runner.py to have a new function _create_evaluators for create all evaluators. Thus we do not need to override the whole _do_test function in the runner which need to add the visualization evaluator of all ranks.

(Note: this ignores all push blocking failures!)

Reviewed By: YanjunChen329

Differential Revision: D43263543

fbshipit-source-id: eca2259277584819dcc5400d47fa4fb142f2ed9b
      b21607b1
    • Yanghan Wang's avatar
      add type annotations to preserve return type · 31197c3e
      Yanghan Wang authored 
      Summary:
X-link: https://github.com/facebookresearch/mobile-vision/pull/137

Pull Request resolved: https://github.com/facebookresearch/d2go/pull/475

Reviewed By: YanjunChen329

Differential Revision: D42148563

fbshipit-source-id: 76b794988bda7f773a734838c79d2de087d7ce94
      31197c3e
  19. 14 Feb, 2023 3 commits
    • Fei Sun's avatar
      Add NUMA binding · 07ddd262
      Fei Sun authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/472

Add NUMA binding to d2go. It equally distributes the GPUs to the CPU sockets so that the CPU traffic, GPU to CPU traffic are all balanced. It helps the diffusion model training, but it is a general technique that can be applied to all models. We still want to manually enable it in each case though, until we are confident that it gives better performance and set it as a default.

NUMA binding is based on jspark1105's work D42827082. Full credit goes to him.

This diff does not enable the feature.

Reviewed By: newstzpz

Differential Revision: D43036817

fbshipit-source-id: fe67fd656ed3980f04bc81909cae7ba2527346fd
      07ddd262
    • Fei Sun's avatar
      Add option to use fused adamw optimizer · 8bb24bb0
      Fei Sun authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/471

adamw recently added an option to use a fused optimizer. It may give better performance than foreach argument. However, we cannot enable it by default, since it requires all parameters to be in CUDA and maybe some other restrictions. So, enable it on a per project basis.

On DALLE2, it results about 23ms faster.

Reviewed By: newstzpz

Differential Revision: D43027327

fbshipit-source-id: 82c6855116094e86386ad2edeea3a74f9e555174
      8bb24bb0
    • Fei Sun's avatar
      Ignore modules · 7ef9d897
      Fei Sun authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/470

Enable ignore FSDP modules. Those modules will not be put in FSDP. It is useful in the diffusion model, where the CLIP model is not used in training. Thus, it is OK to have a separate copy in each GPU. It reduces the CLIP execution time from 63ms to 48ms (15ms reduction). This is mostly because it is a CPU bounded module and in each FSDP block, some code is injected. In addition, it also reduces the FSDP all gather time before the CLIP execution from 56ms to 7ms (49ms reduction).

In total, this change may reduce the CLIP runtime from 119ms to 64ms (63ms reduction)

This feature is controlled by this flag:
    IGNORED_MODULES: ["clip_model"]

Reviewed By: newstzpz

Differential Revision: D42910383

fbshipit-source-id: dc4c12254d45ac45d88329feb63a26ec4ae04aef
      7ef9d897
  21. 05 Feb, 2023 1 commit
    • Maayan Frid-Adar's avatar
      Fix TB train visualization · c4c512ce
      Maayan Frid-Adar authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/465

Training visualization was basically activated only for the first training iterations if TRAIN_LOADER_VIS_MAX_IMAGES and TRAIN_LOADER_VIS_WRITE_PERIOD were set to be > 0. because the MAX_IMAGES was taken as the number of samples to log + the allowed number of samples to load overall. So after the first log to TB it was set to 0 and the visualization was not activated for later training steps (ignoring WRITE_PERIOD).

I've added a TRAIN_LOADER_VIS_MAX_BATCH_IMAGES parameter to set a number of samples to visualize each write period up to the max images defined with TRAIN_LOADER_VIS_MAX_IMAGES

Reviewed By: tglik

Differential Revision: D42832903

fbshipit-source-id: 02a0d9aa4ea6d0ee725120916d26b77843a3e8ab
      c4c512ce
  23. 04 Feb, 2023 1 commit
  25. 03 Feb, 2023 1 commit
  27. 01 Feb, 2023 2 commits
    • Licheng Yu's avatar
      missing keys in _convert_to_d2 · c5bf9222
      Licheng Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/462

Fix errors in `_convert_to_d2`. Sometimes the keys are missing, we don't need remove them.

{F860805441}

Reviewed By: newstzpz

Differential Revision: D42929485

fbshipit-source-id: 8584879df5a07cbe5a864b4f170eef3d5f34dd6c
      c5bf9222
    • Yanghan Wang's avatar
      Allow specifying extra lightning trainer params via `_DEFAULTS_` in yaml · 6940fa9c
      Yanghan Wang authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/461

There're needs for extending trainer parameters that are not in (or conflict with) the base d2go config, this diff adds a way to inject those configs without touching the base d2go config.
- In `get_trainer_params`, it simply checks the `LIGHTNING_TRAINER` and use whatever configs under it.
- Adds `GeneralizedRCNNTaskNoDefaultConfig`, which allows specify default config via yaml file for `GeneralizedRCNNTask`. (also make some changes for prerequisite)
- (next diff) User can add their own config updater by registering it in `CONFIG_UPDATER_REGISTRY`.

Differential Revision: D42928992

fbshipit-source-id: f2a1d8a3f2bec9908bb1af03928611d963b92c0e
      6940fa9c
  29. 23 Jan, 2023 1 commit
  31. 16 Jan, 2023 1 commit
  33. 14 Jan, 2023 1 commit
  35. 13 Jan, 2023 4 commits
    • Anthony Chen's avatar
      Make AMP compatible with FSDP · abf0ca0c
      Anthony Chen authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/458

Make AMP compatible with FSDP. FSDP does not depend on the torch AMP module and implements its own MixedPrecision module. This MixedPrecision module directly saves additional copy of weights in lower precision and use run these tensors in mixed precision training. This is very different from AMP, which automatically casts tensors to lower precision upon tensor operations.

This diff solves some compatibility bugs between AMP and FSDP with 2 changes:
1. Use "never_wrap_policy" as the default dummy autowrap policy.
FSDP Mixed Precision doesn't work with Batchnorm layers. This is because FSDP and other resources like NVidia apex highly discourage running lower precision for batchnorm: https://github.com/pytorch/pytorch/issues/75478. We need to use some autowrap policy in order to let FSDP surpass batchnorm layers in constructing mixed precision.
2. Wrap FSDPWrapper.forward() with autocast()
FSDP Mixed Precision uses lower-precision tensors in computation, which could raise type mismatch error when amp.autocast() is not enabled, like in eval. Thus, we wrap FSDP forward() with autocast()

Reviewed By: wat3rBro

Differential Revision: D41328834

fbshipit-source-id: 18cf94c4ad8d9422ffd3bb335873cd29ac987ae9
      abf0ca0c
    • Anthony Chen's avatar
      Convert local checkpoint to global one automatically in d2go FSDP checkpointer · 5ad2d57e
      Anthony Chen authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/446

## Design
Following D41861308, local checkpoints need to be converted to global ones before  being loaded and used in non-FSDP wrapped models. This diff implements such conversion in d2go checkpointer level to allow automatic conversion with minimal user interference and no new config key.

In previous diff, `FSDPWrapper` has 2 loading modes and 2 saving modes: it uses `load_local_state_dict` to determine whether the ckpt we want to load is local or global, and uses `use_local_state_dict` to decide whether to save new ckpts as local or global. Thus, there are 4 combinations of loading/saving modes:
1. load local + save local
2. load local + save global
3. load global + save local
4. load global + save global

And the local-to-global checkpoint conversion maps to mode 2: load local + save global. Thus, when the checkpointer is in mode 2, it automatically saves the model to a global ckpt right after it loads the local ckpt. Because this happens in checkpointer level, normal training/eval can resume after ckpt conversion. This gives users a consistent and seamless experience with normal training/eval, while also providing a separate ckpt conversion feature via eval-only.

## Usage
Suppose we want to convert local checkpoint `/tmp/model_final`, user can run the same training command with extra args: `MODEL.WEIGHTS=/tmp/model_final` and `FSDP.USE_LOCAL_STATE_DICT=False`

Wiki: https://www.internalfb.com/intern/wiki/Mobile_Vision/Detectron2Go/D2 (https://github.com/facebookresearch/d2go/commit/87374efb134e539090e0b5c476809dc35bf6aedb)Go_Tutorials/Diffusion_Pipeline/Diffusion_Model_Inference/#using-checkpoints-traine

Reviewed By: wat3rBro

Differential Revision: D41926662

fbshipit-source-id: 18a62607a79b0e917d929e9ea85ac1658fb895ca
      5ad2d57e
    • Anthony Chen's avatar
      Support local state dict checkpointing for FSDP · eea6339f
      Anthony Chen authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/457

## Context:

The Pytorch FSDP (Fully Sharded Data Parallel) backend supports two checkpointing modes. The first one is full_state_dict mode, where each FSDP worker summons parameters from other workers to produce a global state dict that can be loaded by non-FSDP models. This mode is the desired mode for checkpointing because checkpoint structures and key names follows the default convention. It's already supported in D39228316 (https://github.com/facebookresearch/d2go/commit/02625ff83207b836df349eadc4a61eb3d4a5810c)

However, when the model is too large to fit into a single GPU memory, this approach would fail because a worker's GPU can't hold all the summoned parameters during checkpoint saving. The rescue is to use the second checkpointing mode: local_state_dict. This mode saves the sharded parameters in each GPU process locally. It can only be loaded by FSDP-wrapped models with the same distributed training settings (i.e. num processes), but it reduces the need for summoning parameters and greatly saves peak GPU memory during training

This diff enables local state dict checkpointing in d2go.

## API:

This diff supports both **saving** local state and **loading** state dict that is locally sharded. Whether to save local state is controlled by `FSDP.USE_LOCAL_STATE`. If `FSDP.USE_LOCAL_STATE=True` and we want to save `output/model_0000001.pth` as in the old pattern, the local checkpoints will be saved as:
```
- output
    - model_0000001
        - rank0.pth
        - rank1.pth
        - rank2.pth
        - rank3.pth
```
Whether to load local state, on the other hand, is controlled by the path of the checkpoint to load. If the path is a file, i.e. `output/model_final.pth`, the file will be loaded as a full state dict by all GPU processes like before. If the path is a directory, i.e. `output/model_final`, the checkpointer will attempt to load `output/model_final/rankX.pth` for rank X.

This API design enables the full combinations of loading local/full states and saving local/full states.

## Conversion to full state dict [Temporary]

Conversion from local state dict to full state dict is needed during an e2e workflow. This will be implemented in another diff

Reviewed By: wat3rBro

Differential Revision: D41861308

fbshipit-source-id: 2e01b601683d06b46f0c5517c6cff30bbcffa8f7
      eea6339f
    • Anthony Chen's avatar
      Rewrite FSDP wrapping as modeling hook · dc6fac12
      Anthony Chen authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/440

Move FSDP wrapping to runner.build_model by rewriting it as a modeling hook

**Motivation**
When a model is too large to run inference on a single GPU, it requires using FSDP with local checkpointing mode to save peak GPU memory. However, in eval_pytorch workflow (train_net with eval-only), models are evaluated without being wrapped by FSDP. This may cause OOM errors for the reasons above. Thus, it may be a better practice to wrap model with FSDP during `runner.build_model(cfg)`, so evaluation can also be run in the same FSDP setting as in training.

This diff moves FSDP wrapping to `runner.build_model(cfg)` by rewriting it as a modeling hook.

**API changes**
* Users need to append `"FSDPModelingHook"` to `MODEL.MODELING_HOOKS` to enable FSDP.
* `FSDP.ALGORITHM` can only be `full` or `grad_optim`

**Note**
It's not possible to unwrap an FSDP model back to the normal model, so FSDPModelingHook.unapply() can't be implemented

Reviewed By: wat3rBro

Differential Revision: D41416917

fbshipit-source-id: f3fc72d574cc6ccbe0d238e48c575926ba5b4d06
      dc6fac12
  37. 10 Jan, 2023 2 commits
  39. 05 Jan, 2023 3 commits
  41. 04 Jan, 2023 1 commit
    • Yanghan Wang's avatar
      upgrade pytorch-lightning version to 1.8.6 · 9e93852d
      Yanghan Wang authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/453

Previous diffs updated the LRScheduler to public version (eg. https://github.com/facebookresearch/detectron2/pull/4709), this also requires newer version of pytorch-lightning. This diff upgrades the lightning version to 1.8.6, also fixes some deprecated call sites of old lightning versions.
- `deepcopy` seems to be supported now, remove `_deepcopy` (there's now not allowed to access `trainer` attributed when it is `None`)
- `dataloader_idx` is removed from `on_train_batch_start`.
- stop using `_accelerator_connector` (the AcceleratorConnector doesn't have those attributes anymore).
- deprecated `on_pretrain_routine_end` -> `on_fit_start`

Reviewed By: YanjunChen329

Differential Revision: D42319019

fbshipit-source-id: ba46abbd98da96783e15d187a361fda47dc7d4d6
      9e93852d
  43. 20 Dec, 2022 1 commit
  45. 19 Dec, 2022 1 commit