Skip to content

GitLab

Switch branch/tag
  1. 14 Jan, 2023 1 commit
  3. 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
  5. 10 Jan, 2023 2 commits
  7. 05 Jan, 2023 3 commits
  9. 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
  11. 20 Dec, 2022 1 commit
  13. 19 Dec, 2022 4 commits
    • Yanghan Wang's avatar
      add check to `import_runner` · fb41d071
      Yanghan Wang authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/450

move checking logic into `import_runner`, simplfy `_is_lightning_task`

Reviewed By: mcimpoi

Differential Revision: D42105853

fbshipit-source-id: 5fd51865a01f2cbac38aaedcac49207c26172ab9
      fb41d071
    • Haroun Habeeb's avatar
      temp_new_allowed · 7bed2910
      Haroun Habeeb authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/438

Adding new fields to a config is only allowed if `new_allowed=True`. yacs `CfgNode` provides a `set_new_allowed(value: bool)` function.

We create a context manager like `temp_defrost` but for new_allowed to use it. We also implement unit test for the same

Reviewed By: yanglinfang, newstzpz, wat3rBro

Differential Revision: D41748992

fbshipit-source-id: 71d048511476001ca96e6b36dde4d177b11268d7
      7bed2910
    • Yanghan Wang's avatar
      separate TestNetOutput and TrainNetOutput · e2537c82
      Yanghan Wang authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/449

separate TestNetOutput and TrainNetOutput
- update d2go binaries
- update operators / workflows

Reviewed By: mcimpoi

Differential Revision: D42103714

fbshipit-source-id: 53f318c79d7339fb6fcfc3486e8b9cf249a598bf
      e2537c82
    • Anton Rigner's avatar
      Fix WeightedSampler to also work with adhoc datasets · ab49d0b6
      Anton Rigner authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/437

# Problem
- We use `TRAIN_CATEGORIES` to overrider the classes for convenient experimentation, to not have to re-map the JSON file
- But it's not possible to use the WeightedTrainingSampler with specified repeat factors (`DATASETS.TRAIN_REPEAT_FACTOR`) when also overriding the classes to use for training (ad-hoc datasets), because the underlying dataset name doesn't match the datasets specified in the `TRAIN_REPEAT_FACTOR` pairs (mapping between <dataset_name, repeat_factor>)

# Fix

- Update the dataset names for the REPEAT_FACTORS mapping as well, if we have enabled the `WeightedTrainingSampler` and use ad-hoc datasets.

Reviewed By: wat3rBro

Differential Revision: D41765638

fbshipit-source-id: 51dad484e4d715d2de900b5d0b7c7caa19903fb7
      ab49d0b6
  15. 16 Dec, 2022 1 commit
  17. 12 Dec, 2022 2 commits
  19. 09 Dec, 2022 1 commit
  21. 08 Dec, 2022 1 commit
  23. 30 Nov, 2022 6 commits
    • Matthew Yu's avatar
      support caching tuples · dece58ba
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/432

We support caching of tuples since they behave similarly to lists

Reviewed By: XiaoliangDai

Differential Revision: D41483876

fbshipit-source-id: 9d741074f8e2335ddd737ae3f1bdb288910f5564
      dece58ba
    • Matthew Yu's avatar
      algorithm · 150db2d1
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/431

Add a generic domain adaptation algorithm. This algorithm:
* gets domain0 data out of the dataloader
* runs domain0 data into the model and saves target layer output
* gets domain1 data of the dataloader
* runs domain1 data into the model and saves target layer output
* runs domain adaptation loss on domain0, domain1 outputs
* combines losses using model training iteration

This diffs adds `get_preprocess_domain0_input` and `get_preprocess_domain1_input` to the distillation helper. These are functions that the user can use to convert the dataloader output to something that will be used by the model (e.g., pull the domain0 or domain1 key out of a dataloader that returns a dict).

Differential Revision: D40970724

fbshipit-source-id: fff050fbe864654fa6cb0df927f6843855ec1c14
      150db2d1
    • Matthew Yu's avatar
      support registering layer losses to model · c4860c5b
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/430

We add losses in distillation by instantiating them in the distillation algorithm's init and then running them during the forward pass.

However this has some issues:
* the losses are not registered as a module in the model since they we organize them as a list of layerlossmetadata => this means that things like AMP do not behave as expected
* the losses are not on the same device as the rest of the model since they are created potentially after the model is moved to a new device

This diff solves both of these issues by including a helper function that registers and moves the losses to the same device as the model. `register_layer_losses_and_to_device` takes as input `List[LayerLossMetadata]`, moves the losses to the same device as the model and then registers these losses to the model.

Differential Revision: D41296932

fbshipit-source-id: ae7ae0847bce1b5cc481d838b9cae69cea424f25
      c4860c5b
    • Matthew Yu's avatar
      support ignoring teacher · 909de50d
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/429

Add a teacher type called `no_teacher` which can be specified by the user in the case they ignore the teacher (e.g., domain adaptation). Building the teacher just returns a noop (`nn.Identity`)

Differential Revision: D40971788

fbshipit-source-id: fc49ac44224c92806a7be253eefb8454305814eb
      909de50d
    • Peizhao Zhang's avatar
      add an augmentation to pad image to square. · c2d7dbab
      Peizhao Zhang authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/428

add an augmentation to pad image to square.
* For example, image with shape (10, 7, 3) will become (10, 10, 3) and pad with value specified by `pad_value`.

Reviewed By: tax313, wat3rBro

Differential Revision: D41545182

fbshipit-source-id: 6d5fd9d16984a9904d44f22386920cdf130edda7
      c2d7dbab
    • Matthew Yu's avatar
      set cache in recorded layers · 30ac5858
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/433

Distillation uses a module called `CachedLayer` to record the outputs of a layer to an internal dict. This dict is typically initialized by the object itself and any value is overwritten every time the model runs.

However, sometimes we need more than one output run of the layer (e.g., domain adaptation => we run the model on real, then synthetic data and need to use both outputs).

This diff adds a helper to set externally set the cache dict of a model. In other words, we can run `set_cache_dict` on some model to change the dict used by all `CachedLayer` in the model. This allows us to run the model and record some outputs, then change the cache dict and rerun the model to save different outputs.

Differential Revision: D40970577

fbshipit-source-id: 49cb851af49ae193d0c8ac9218e02fdaf4e6587b
      30ac5858
  25. 28 Nov, 2022 1 commit
  27. 23 Nov, 2022 2 commits
  29. 22 Nov, 2022 1 commit
    • Matthew Yu's avatar
      add default layer losses and loss combiner · 419974bb
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/421

Add some reasonable defaults when running knowledge distillation
* get_default_kd_image_classification_layer_losses => returns cross entropy loss on the output of the student classification layer and the teacher output (this is what the imagenet distillation uses)
* DefaultLossCombiner => simple function to multiply the losses by some weights

Unsure if these should go in `distillation.py` or a separate place (e.g., defaults or classification)

Reviewed By: chihyaoma

Differential Revision: D40330718

fbshipit-source-id: 5887566d88e3a96d01aca133c51041126b2692cc
      419974bb
  31. 21 Nov, 2022 1 commit
    • Yanghan Wang's avatar
      add configure_dataset_creation for lightning · 0ea6bc1b
      Yanghan Wang authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/423

The `DefaultTask` has forked the implementation of some runner methods from `Detectron2GoRunner`, which is not necessary since there should be no differences. This might cause issue that we update the one from `Detectron2GoRunner` but forgot about `DefaultTask`.

Reviewed By: chihyaoma

Differential Revision: D41350485

fbshipit-source-id: 38a1764a7cc77dc13939ac7d59f35584bf9dab9b
      0ea6bc1b
  33. 19 Nov, 2022 1 commit
    • Matthew Yu's avatar
      kd algorithm · 9ec4f2bf
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/420

Adds knowledge distillation as a generic algorithm that can be used by various projects.

If eval, the algorithm just returns the result of the student model.

If training, the algorithm feeds the input into both the student and teacher model. The user provides a list of `LayerLossMetadata` that provides the layers and losses run on these layers. The algorithm uses dynamic mixin to record the outputs of the relevant layers and compute the losses after both models are run.

We provide student and teacher preprocessing as a placeholder before we support a more generic dataloader which can provide different inputs to the student and teacher (e.g., as of now, if you want to provide the teacher with a larger input then the dataloader should return a large input and the student preprocessing can downsample the input).

We add the following functions as part of the user customizable distillation helper:
* get_teacher => return a teacher that can be used directly by the KD algorithm
* get_layer_losses => return a list of `LayerLossMetadata` that provides the layers and losses
* get_preprocess_student_input => manipulate the output of the dataloader before passing to the student
* get_preprocess_teacher_input => manipulate the output of the dataloader before passing to the teacher
* get_combine_losses => since we may want to weight the student and distillation losses, return a function that can manipulate the loss_dict

Reviewed By: chihyaoma

Differential Revision: D40326412

fbshipit-source-id: 2fb0e818a7d5b120d62fb7aba314ff96cc7e10c5
      9ec4f2bf
  35. 18 Nov, 2022 1 commit
  37. 17 Nov, 2022 3 commits
    • Anthony Chen's avatar
      Integrate PyTorch Fully Sharded Data Parallel (FSDP) · 02625ff8
      Anthony Chen authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/396

Integrate PyTorch FSDP, which supports two sharding modes: 1. gradient + optimizer sharding; 2. full model sharding (params + gradient + optimizer). This feature is enabled in the train_net.py code path.

Sources
* Integration follows this tutorial: https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html

API changes
* Add new config keys to support the new feature. Refer to mobile-vision/d2go/d2go/trainer/fsdp.py for the full list of config options
* Add `FSDPCheckpointer` as an inheritance of `QATCheckpointer` to support special loading/saving logic for FSDP models

Reviewed By: wat3rBro

Differential Revision: D39228316

fbshipit-source-id: 342ecb3bcbce748453c3fba2d6e1b7b7e478473c
      02625ff8
    • Matthew Yu's avatar
      add class to keep track of loss metadata and function to compute losses · 0316fed4
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/419

This diff adds a metadata class `LayerLossMetadata` to help keep track of the losses we want to compute over layers. The class contains the type of loss, loss name, and layer names.

This diff adds a helper function to iterate over a list of `LayerLossMetadata` and return a dict containing the results.

Reviewed By: chihyaoma

Differential Revision: D40286564

fbshipit-source-id: b269dc63cc90a437ca279379d759c3106016327c
      0316fed4
    • Matthew Yu's avatar
      add a helper to record layers in a model · 53c4c2c1
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/418

This diff adds a function that can be used to add `CachedLayers` to a model. Function iterates over named modules and dynamically mixes in `CachedLayer` to target modules.

This diff adds a function to remove the cached layers.

Reviewed By: Minione

Differential Revision: D40285806

fbshipit-source-id: 3137d19927d8fb9ec924a77c9085aea29fe94d5e
      53c4c2c1
  39. 16 Nov, 2022 2 commits
    • Matthew Yu's avatar
      support a layer that saves outputs · 120b463c
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/417

This diff adds a layer `CachedLayer` which is meant to be used with dynamic mixin. This layer runs the original module and clones the output into a dictionary provided by the user.

The main use case is in distillation where we dynamically mixin these layers to the layers that the user wants to compute various losses.

See subsequent diffs to get integration with distillation.

Reviewed By: Minione

Differential Revision: D40285573

fbshipit-source-id: 2058deff8b96f63aebd1e9b9933a5352b5197111
      120b463c
    • Matthew Yu's avatar
      update teacher to support models where device is a property · 0f27e90f
      Matthew Yu authored 
      Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/416

Distillation assumes teacher model has an attribute "device". Sometimes this attribute is actually a property (e.g., generalizedrcnn) but there is zero guarantee that it exists. We add a helper function to move the model to the device and add this attribute if needed.

Reviewed By: chihyaoma

Differential Revision: D40283954

fbshipit-source-id: 42921653eac8a79499e22edac29aa6aeac016e8a
      0f27e90f
  41. 15 Nov, 2022 1 commit