Skip to content
  • wayi1's avatar
    [Model Averaging] Support hierarchical model averaging (#73285) · 0bb3b065
    wayi1 authored
    Summary:
    Implement hierarchical model averaging proposed in https://github.com/pytorch/pytorch/issues/71325.
    
    Unit tests are added. Since I don't have access to 4-GPU machines in open-source environment, expect that the branch with the prefix of `ci-all` can run the test that requires 4 GPUs.
    
    In the future, the internals of `PeriodicModelAveraging` can be simplified as an implementation of a specialized hierarchical model averaging, where `period_group_size_dict` only has a pair of period and world size.
    
    Pull Request resolved: https://github.com/pytorch/pytorch/pull/73285
    
    Reviewed By: mrshenli
    
    Differential Revision: D34457792
    
    Pulled By: rohan-varma
    
    fbshipit-source-id: 39a6c5bf8a2852b6394a56abbad17b8a909b9fba
    (cherry picked from commit 5f543d46)
    0bb3b065
This project is licensed under the BSD 3-Clause "New" or "Revised" License. Learn more