As of now, I plan to measure the time taken for gradient accumulation

Are you referring to the AccumulateGrad 4 function? If so, the autograd profiler would display time taken for this function.

and the number of bits communicated (for each iteration)

This should be possible in the current proposal of the communication hook since you can add up the bits for all the buckets.

I might need to even find the bits communicated for each layer in the future to explore layer wise compression.

This is probably something that you can’t do with the existing hook since it provides the entire bucket to the user and currently there is no way to split out individual parameters from the bucket.