Add documentation

Project.toml

@@ -62,6 +62,7 @@ ADTypes = "0.2"
 Adapt = "4"
 Aqua = "0.8.4"
 ArrayInterface = "7.8"
+CUDA = "5.2"
 ChainRules = "1.62"
 ChainRulesCore = "1.21"
 ComponentArrays = "0.15.11"
@@ -80,13 +81,15 @@ LuxCore = "0.1.12"
 LuxDeviceUtils = "0.1.19"
 LuxLib = "0.3.11"
 LuxTestUtils = "0.1.15"
-MLUtils = "0.4"
+MLUtils = "0.4.3"
+MPI = "0.20.19"
 MacroTools = "0.5.13"
 Markdown = "1.10"
 NCCL = "0.1.1"
 Optimisers = "0.3"
 Pkg = "1.10"
 PrecompileTools = "1.2"
+Preferences = "1.4.3"
 Random = "1.10"
 ReTestItems = "1.23.1"
 Reexport = "1"

docs/Project.toml

@@ -14,5 +14,6 @@ LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

docs/make.jl

@@ -38,14 +38,16 @@ pages = [
         "manual/freezing_model_parameters.md",
         "manual/gpu_management.md",
         "manual/migrate_from_flux.md",
-        "manual/weight_initializers.md"
+        "manual/weight_initializers.md",
+        "manual/distributed_utils.md"
     ],
     "API Reference" => [
         "Lux" => [
             "api/Lux/layers.md",
             "api/Lux/utilities.md",
             "api/Lux/contrib.md",
-            "api/Lux/switching_frameworks.md"
+            "api/Lux/switching_frameworks.md",
+            "api/Lux/distributed_utils.md",
         ],
         "Accelerator Support" => [
             "api/Accelerator_Support/LuxAMDGPU.md",

docs/src/.vitepress/config.mts

@@ -73,7 +73,8 @@ export default defineConfig({
                             { text: 'Built-In Layers', link: '/api/Lux/layers' },
                             { text: 'Utilities', link: '/api/Lux/utilities' },
                             { text: 'Experimental', link: '/api/Lux/contrib' },
-                            { text: 'InterOp', link: '/api/Lux/switching_frameworks' }
+                            { text: 'InterOp', link: '/api/Lux/switching_frameworks' },
+                            { text: 'DistributedUtils', link: '/api/Lux/distributed_utils' }
                         ]
                     },
                     {
@@ -146,7 +147,8 @@ export default defineConfig({
                     { text: 'Freezing Model Parameters', link: '/manual/freezing_model_parameters' },
                     { text: 'GPU Management', link: '/manual/gpu_management' },
                     { text: 'Migrating from Flux to Lux', link: '/manual/migrate_from_flux' },
-                    { text: 'Initializing Weights', link: '/manual/weight_initializers' }]
+                    { text: 'Initializing Weights', link: '/manual/weight_initializers' },
+                    { text: 'Distributed Data Parallel Training', link: '/manual/distributed_utils' },]
             },
             "/api/": {
                 text: 'API Reference', collapsed: false, items: [
@@ -155,7 +157,8 @@ export default defineConfig({
                             { text: 'Built-In Layers', link: '/api/Lux/layers' },
                             { text: 'Utilities', link: '/api/Lux/utilities' },
                             { text: 'Experimental Features', link: '/api/Lux/contrib' },
-                            { text: 'Switching between Deep Learning Frameworks', link: '/api/Lux/switching_frameworks' }]
+                            { text: 'Switching between Deep Learning Frameworks', link: '/api/Lux/switching_frameworks' },
+                            { text: 'DistributedUtils', link: '/api/Lux/distributed_utils' }]
                     },
                     {
                         text: 'Accelerator Support', collapsed: false, items: [

docs/src/api/Lux/distributed_utils.md

@@ -0,0 +1,58 @@
+# Distributed Utils
+
+!!! note
+
+    These functionalities are available via the `Lux.DistributedUtils` module.
+
+```@meta
+CurrentModule = Lux
+```
+
+## Index
+
+```@index
+Pages = ["distributed_utils.md"]
+```
+
+## [Backends](@id communication-backends)
+
+```@docs
+MPIBackend
+NCCLBackend
+```
+
+## Initialization
+
+```@docs
+DistributedUtils.initialize
+DistributedUtils.initialized
+DistributedUtils.get_distributed_backend
+```
+
+## Helper Functions
+
+```@docs
+DistributedUtils.local_rank
+DistributedUtils.total_workers
+```
+
+## Communication Primitives
+
+```@docs
+DistributedUtils.allreduce!
+DistributedUtils.bcast!
+DistributedUtils.reduce!
+DistributedUtils.synchronize!!
+```
+
+## Optimizers.jl Integration
+
+```@docs
+DistributedUtils.DistributedOptimizer
+```
+
+## MLUtils.jl Integration
+
+```@docs
+DistributedUtils.DistributedDataContainer
+```

docs/src/introduction/index.md

@@ -75,6 +75,7 @@ st_opt, ps = Optimisers.update(st_opt, ps, gs)
 using Lux, Random, Optimisers, Zygote
 # using LuxCUDA, LuxAMDGPU, Metal # Optional packages for GPU support
 import Lux.Experimental: @compact
+using Printf  # For pretty printing
 ```
 
 We will define a custom MLP using the `@compact` macro. The macro takes in a list of
@@ -117,7 +118,7 @@ for epoch in 1:1000
         return sum(abs2, y .- y_data), st_
     end
     gs = only(pb((one(loss), nothing)))
-    epoch % 100 == 1 && println("Epoch: $(epoch) | Loss: $(loss)")
+    epoch % 100 == 1 && @printf "Epoch: %04d \t Loss: %10.9g\n" epoch loss
     Optimisers.update!(st_opt, ps, gs)
 end
 ```

docs/src/manual/distributed_utils.md

@@ -0,0 +1,55 @@
+# Distributed Data Parallel Training
+
+!!! tip
+
+    For a fully functional example, see the
+    [ImageNet Training Example](https://github.com/LuxDL/Lux.jl/tree/main/examples/ImageNet)
+
+DDP Training using `Lux.DistributedUtils` is a spiritual successor to
+[FluxMPI.jl](https://github.com/avik-pal/FluxMPI.jl), but has some key differences.
+
+## Guide to Integrating DistributedUtils into your code
+
+## [GPU-Aware MPI](@id gpu-aware-mpi)
+
+If you are using a custom MPI build that supports CUDA or ROCM, you can use the following
+preferences with [Preferences.jl](https://github.com/JuliaPackaging/Preferences.jl):
+
+1. `LuxDistributedMPICUDAAware` - Set this to `true` if your MPI build is CUDA aware.
+2. `LuxDistributedMPIROCMAware` - Set this to `true` if your MPI build is ROCM aware.
+
+By default, both of these values are set to `false`.
+
+## Migration Guide from `FluxMPI.jl`
+
+Let's compare the changes we need to make wrt the
+[FluxMPI.jl integration guide](https://avik-pal.github.io/FluxMPI.jl/dev/guide/).
+
+1. `FluxMPI.Init` is now [`DistributedUtils.initialize`](@ref).
+2. `FluxMPI.synchronize!(x)` needs to be changed to
+   `x_new = DistributedUtils.synchronize!!(backend, x)`.
+3. [`DistributedUtils.DistributedDataContainer`](@ref),
+   [`DistributedUtils.local_rank`](@ref), and
+   [`DistributedUtils.DistributedOptimizer`](@ref) need `backend` as  the first input.
+
+And that's pretty much it!
+
+### Removed Functionality
+
+1. `FluxMPI.allreduce_gradients` no longer exists. Previously this was needed when CUDA
+   communication was flaky, with `NCCL.jl` this is no longer the case.
+2. `FluxMPIFluxModel` has been removed. `DistributedUtils` no longer works with `Flux`.
+
+### Key Differences
+
+1. `FluxMPI.synchronize!` is now `DistributedUtils.synchronize!!` to highlight the fact
+   that some of the inputs are not updated in-place.
+2. All of the functions now require a [communication backend](@ref communication-backends)
+   as input.
+3. We don't automatically determine if the MPI Implementation is CUDA or ROCM aware. See
+   [GPU-aware MPI](@ref gpu-aware-mpi) for more information.
+4. Older [`Lux.gpu`](@ref) implementations used to "just work" with `FluxMPI.jl`. We expect
+   [`gpu_device`](@ref) to continue working as expected, however, we recommend using
+   [`gpu_device`](@ref) after calling [`DistributedUtils.initialize`](@ref) to avoid any
+   mismatch between the device set via `DistributedUtils` and the device stores in
+   `LuxCUDADevice` or `LuxAMDGPUDevice`.

ext/LuxComponentArraysExt.jl

@@ -1,7 +1,7 @@
 module LuxComponentArraysExt
 
 using ComponentArrays: ComponentArrays, ComponentArray, FlatAxis
-using Lux: Lux
+using Lux: Lux, DistributedUtils
 
 # Empty NamedTuple: Hack to avoid breaking precompilation
 function ComponentArrays.ComponentArray(data::Vector{Any}, axes::Tuple{FlatAxis})

ext/LuxFluxExt.jl

@@ -143,8 +143,8 @@ function __from_flux_adaptor(l::Flux.ConvTranspose; preserve_ps_st::Bool=false,
     if preserve_ps_st
         _bias = l.bias isa Bool ? nothing :
                 reshape(copy(l.bias), ntuple(_ -> 1, length(k))..., out_chs, 1)
-        return Lux.ConvTranspose(k, in_chs * groups => out_chs, l.σ; l.stride, pad,
-            l.dilation, groups, use_bias=!(l.bias isa Bool),
+        return Lux.ConvTranspose(k, in_chs * groups => out_chs, l.σ; l.stride,
+            pad, l.dilation, groups, use_bias=!(l.bias isa Bool),
             init_weight=__copy_anonymous_closure(Lux._maybe_flip_conv_weight(l.weight)),
             init_bias=__copy_anonymous_closure(_bias))
     else

ext/LuxMLUtilsExt.jl

@@ -19,4 +19,4 @@ function MLUtils.getobs(dc::DistributedUtils.DistributedDataContainer, idx)
     return MLUtils.getobs(dc.data, dc.idxs[idx])
 end
 
-end
+end

ext/LuxMPINCCLExt.jl

@@ -8,15 +8,15 @@ using Setfield: @set!
 
 function DistributedUtils.__initialize(
         ::Type{NCCLBackend}; cuda_devices=nothing, amdgpu_devices=missing)
-    DistributedUtils.NCCL_Initialized[] = true
     @assert amdgpu_devices===missing "`AMDGPU` is not supported by `NCCL`."
-    DistributedUtils.__initialize(Val(:MPI); cuda_devices, amdgpu_devices)
+    DistributedUtils.__initialize(MPIBackend; cuda_devices, amdgpu_devices)
+    DistributedUtils.NCCL_Initialized[] = true
     return
 end
 
 function DistributedUtils.__get_distributed_backend(::Type{NCCLBackend})
     unique_id = NCCL.UniqueID()  # Generate on all ranks to know the type
-    mpi_backend = DistributedUtils.__get_distributed_backend(Val(:MPI))
+    mpi_backend = DistributedUtils.__get_distributed_backend(MPIBackend)
     buf = [unique_id.internal...]
     DistributedUtils.bcast!(mpi_backend, buf; root=0)
     @set! unique_id.internal = Tuple(buf)

ext/LuxSimpleChainsExt.jl

@@ -6,12 +6,12 @@ using Lux: SimpleChainsModelConversionError, __to_simplechains_adaptor,
            __fix_input_dims_simplechain
 using Random: AbstractRNG
 
-function __fix_input_dims_simplechain(layers::Vector, input_dims)
+function Lux.__fix_input_dims_simplechain(layers::Vector, input_dims)
     L = Tuple(layers)
     return SimpleChains.SimpleChain{typeof(input_dims), typeof(L)}(input_dims, L)
 end
 
-function __fix_input_dims_simplechain(layers, input_dims)
+function Lux.__fix_input_dims_simplechain(layers, input_dims)
     @warn "The model provided is not a `Chain`. Trying to wrap it into a `Chain` but this \
            might fail. Please consider using `Chain` directly (potentially with \
            `disable_optimizations = true`)."
@@ -21,16 +21,16 @@ end
 __equivalent_simplechains_fn(::typeof(Lux.relu)) = SimpleChains.relu
 __equivalent_simplechains_fn(f::F) where {F} = f
 
-function __to_simplechains_adaptor(layer::Dense{use_bias}) where {use_bias}
+function Lux.__to_simplechains_adaptor(layer::Dense{use_bias}) where {use_bias}
     return SimpleChains.TurboDense{use_bias}(
         __equivalent_simplechains_fn(layer.activation), layer.out_dims)
 end
 
-function __to_simplechains_adaptor(layer::Chain)
+function Lux.__to_simplechains_adaptor(layer::Chain)
     return reduce(vcat, map(__to_simplechains_adaptor, layer.layers))
 end
 
-function __to_simplechains_adaptor(layer::Conv)
+function Lux.__to_simplechains_adaptor(layer::Conv)
     if all(==(1), layer.stride) &&
        layer.groups == 1 &&
        all(==(1), layer.dilation) &&
@@ -42,13 +42,13 @@ function __to_simplechains_adaptor(layer::Conv)
                                             supported."))
 end
 
-function __to_simplechains_adaptor(layer::Dropout)
+function Lux.__to_simplechains_adaptor(layer::Dropout)
     layer.dims isa Colon && return SimpleChains.Dropout(layer.p)
     throw(SimpleChainsModelConversionError("Dropout with non-standard parameters not \
                                             supported."))
 end
 
-function __to_simplechains_adaptor(layer::FlattenLayer)
+function Lux.__to_simplechains_adaptor(layer::FlattenLayer)
     if layer.N === nothing
         throw(SimpleChainsModelConversionError("`FlattenLayer(nothing)` not supported. For \
                                                 `SimpleChains.Flatten` you must use \
@@ -57,15 +57,15 @@ function __to_simplechains_adaptor(layer::FlattenLayer)
     return SimpleChains.Flatten(layer.N)
 end
 
-function __to_simplechains_adaptor(layer::MaxPool)
+function Lux.__to_simplechains_adaptor(layer::MaxPool)
     if layer.stride == layer.k && (!(layer.pad isa SamePad) && all(==(0), layer.pad))
         return SimpleChains.MaxPool(layer.k)
     end
     throw(SimpleChainsModelConversionError("MaxPool with non-standard parameters not \
                                             supported."))
 end
 
-__to_simplechains_adaptor(layer) = throw(SimpleChainsModelConversionError(layer))
+Lux.__to_simplechains_adaptor(layer) = throw(SimpleChainsModelConversionError(layer))
 
 function Lux.initialparameters(rng::AbstractRNG, layer::SimpleChainsLayer)
     return (; params=Array(SimpleChains.init_params(layer.layer; rng)))

src/Lux.jl

@@ -25,7 +25,7 @@ using PrecompileTools: @recompile_invalidations
                     initialparameters, initialstates, parameterlength, statelength,
                     inputsize, outputsize, update_state, trainmode, testmode, setup, apply,
                     display_name, replicate
-    import LuxDeviceUtils: get_device, set_device!
+    using LuxDeviceUtils: get_device
 end
 
 @reexport using LuxCore, LuxLib, LuxDeviceUtils, WeightInitializers

src/distributed/backend.jl

@@ -4,7 +4,7 @@ abstract type AbstractLuxDistributedBackend end
     MPIBackend(comm = nothing)
 
 Create an MPI backend for distributed training. Users should not use this function directly.
-Instead use [`DistributedUtils.get_distributed_backend(Val(:NCCL))`](@ref).
+Instead use [`DistributedUtils.get_distributed_backend(MPIBackend)`](@ref).
 """
 struct MPIBackend{C} <: AbstractLuxDistributedBackend
     comm::C
@@ -21,7 +21,7 @@ end
     NCCLBackend(comm = nothing, mpi_backend = nothing)
 
 Create an NCCL backend for distributed training. Users should not use this function
-directly. Instead use [`DistributedUtils.get_distributed_backend(Val(:NCCL))`](@ref).
+directly. Instead use [`DistributedUtils.get_distributed_backend(NCCLBackend)`](@ref).
 """
 struct NCCLBackend{C, M <: Union{Nothing, MPIBackend}} <: AbstractLuxDistributedBackend
     comm::C

src/distributed/public_api.jl

@@ -220,6 +220,10 @@ end
 `data` must be compatible with `MLUtils` interface. The returned container is compatible
 with `MLUtils` interface and is used to partition the dataset across the available
 processes.
+
+!!! danger
+
+    `MLUtils.jl` must be installed and loaded before using this.
 """
 @concrete struct DistributedDataContainer
     data
@@ -250,6 +254,10 @@ averages the gradients across the processes using Allreduce.
 ## Arguments
 
   - `optimizer`: An Optimizer compatible with the Optimisers.jl package
+
+!!! danger
+
+    `Optimisers.jl` must be installed and loaded before using this.
 """
 function DistributedOptimizer(backend::AbstractLuxDistributedBackend, opt)
     mod = Base.get_extension(@__MODULE__, :LuxOptimisersExt)