[draft] support async param gather in layer-wise optimizer

[FDecaYed]

What does this PR do ?

This is draft code now, so likely doesn't run. It's meant for demo how to support the feature before proper implementation.

Current Architecture

DistributedOptimizer implements async param gather with these components:

• Bucket-based param organization - Parameters are grouped into buckets within _ParamAndGradBucketGroup
• start_param_sync() - Launches async all-gather using torch.distributed.all_gather with async_op=True
• finish_param_sync() - Waits for completion and dispatches the next bucket's all-gather
• Forward pre-hooks - Registered on modules to call finish_param_sync() when a module needs its params
• assume even sharding and allgather directly into tensor(param buffer)

Changes needed for layer-wise

• let each bucket know which param is on which rank, and group them into lists(lw_params_list in addition to params_list that each bucket already holds, think it as per-bucket layerwise sharding)
• change all_gather_into_tensor to all_gather(_v)
• copy from all_gatherv output into param inside finish_param_sync()

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[mkhona-nvidia]

does async allgather with layerwise still require use-distributed-optimizer? I think yes, to let DDP make the buckets right?

[FDecaYed]

for simplicity to demo the idea, I didn't touch that part. currently if use-distributed-optimizer is off, then all async related functionality will be turned off and code errors out.

But technically this is not required, we just need to change those check from if use-distributed-optimizer to if (use-distributed-optimizer or use-layer-wise)

[mkhona-nvidia]

Another feature is that the current layerwise allows reshardable EP sizes in the middle of training. AFAIK, using adam's distopt does not allow this because of the way the DDP buckets are constructed. Can we avoid this and continue to allow reshardable EP sizes in the middle of training with layerwise?

[skyw]

Reuse "overlap_param_gather" if the flag indicates the same thing, no need to introduce new names.

[FDecaYed]

Another feature is that the current layerwise allows reshardable EP sizes in the middle of training. AFAIK, using adam's distopt does not allow this because of the way the DDP buckets are constructed. Can we avoid this and continue to allow reshardable EP sizes in the middle of training with layerwise?

my rough feeling is it should still be supported. the async feature is just moving param allgather from within optimizer.step into forward pre-hook calling bucket function. Checkpointing should not be affected

[mkhona-nvidia]

Another feature is that the current layerwise allows reshardable EP sizes in the middle of training. AFAIK, using adam's distopt does not allow this because of the way the DDP buckets are constructed. Can we avoid this and continue to allow reshardable EP sizes in the middle of training with layerwise?

my rough feeling is it should still be supported. the async feature is just moving param allgather from within optimizer.step into forward pre-hook calling bucket function. Checkpointing should not be affected

I guess the issue is that we don't have a good understanding of why adamW distopt cannot allow resharding of EP in the middle of training. My feeling is that it is because of the way the DDP buckets are made (separate bucketing scheme for EP=true and separate bucketing scheme for rest of the network). If the layerwise distopt overlap_param_gather uses the same scheme again, I think we will run into the same issue

[FDecaYed]

Another feature is that the current layerwise allows reshardable EP sizes in the middle of training. AFAIK, using adam's distopt does not allow this because of the way the DDP buckets are constructed. Can we avoid this and continue to allow reshardable EP sizes in the middle of training with layerwise?

my rough feeling is it should still be supported. the async feature is just moving param allgather from within optimizer.step into forward pre-hook calling bucket function. Checkpointing should not be affected

I guess the issue is that we don't have a good understanding of why adamW distopt cannot allow resharding of EP in the middle of training. My feeling is that it is because of the way the DDP buckets are made (separate bucketing scheme for EP=true and separate bucketing scheme for rest of the network). If the layerwise distopt overlap_param_gather uses the same scheme again, I think we will run into the same issue

I think the EP buckets are separate from regular DP weight no matter if dist-opt is on or not. and async-gather doesn't affect it either.
Here is my guess why dist-opt doesn't support change EP(haven't check code recently, from memory so guess):
there are 2 view of the param buffer, the buffer view(every thing flattened) and the param view(list of param, with their .data backed by the cont buffer). the way dist-opt support torch_dist ckpting is just save the 'buffer view' as it is, which is tightly coupled with EP and cannot be changed. but layerwise safe by each parameter(param view), thus is not affected.
and reason dist-opt have to save it that way is each master parameter/state can be sharded to multiple rank. this is not a problem for layerwise