Does FSDP v2 have the best performance?

[dz1iang]

Hi, when I set fsdp_reshard_after_forward: False, the training speed increased by approximately 5-7%（tokens_per_second_per_gpu）. Are there any other configurations that affect performance? Or where do you recommend referring to for configurations?

In addition, the setting of gradient_accumulation_steps does not affect the speed. Generally speaking, setting a larger value will reduce the frequency of communication and speed up the training. The model used in the experiment is Qwen 2.5 3B.

[ebsmothers]

Hi @dz1iang thanks for creating the issue. Regarding resharding after forward -- in general I would expect this. Setting reshard_after_forward=True means we are doing some additional comms to reduce the overall memory. So if you have extra memory to spare and want to train faster, then reshard_after_forward=False should help (as you observed). There is also some discussion of this in section 5.4 of the original FSDP paper. If you are looking for other settings to experiment with, I would recommend both the fully_shard API reference and torchtitan's doc page comparing FSDP1 and FSDP2. Of course you can also experiment with different sharding strategies -- our default is to shard every transformer layer but you can plug into this API if you want to do something more custom.

I am not sure I fully understand your comment about gradient accumulation steps. Is your point that you would expect to see a speedup by enabling it, but you don't observe one in practice? If so, one thing to note is that we do sync the gradients on each backward pass. Again, this is trading off comms and memory though.. it is possible to disable gradient sync on all but the last iteration, but you then need to keep a full copy of the grads (and the loss calculation becomes a bit trickier).

[dz1iang]

Hi @dz1iang thanks for creating the issue. Regarding resharding after forward -- in general I would expect this. Setting reshard_after_forward=True means we are doing some additional comms to reduce the overall memory. So if you have extra memory to spare and want to train faster, then reshard_after_forward=False should help (as you observed). There is also some discussion of this in section 5.4 of the original FSDP paper. If you are looking for other settings to experiment with, I would recommend both the fully_shard API reference and torchtitan's doc page comparing FSDP1 and FSDP2. Of course you can also experiment with different sharding strategies -- our default is to shard every transformer layer but you can plug into this API if you want to do something more custom.

The straightforward comparison is that there is currently no way to support ZeRO-1.

[dz1iang]

I am not sure I fully understand your comment about gradient accumulation steps. Is your point that you would expect to see a speedup by enabling it, but you don't observe one in practice? If so, one thing to note is that we do sync the gradients on each backward pass. Again, this is trading off comms and memory though.. it is possible to disable gradient sync on all but the last iteration, but you then need to keep a full copy of the grads (and the loss calculation becomes a bit trickier).

Yes, if we want to achieve such an improvement, are there any specific reference cases?

[nathan-az]

@ebsmothers I was looking for ways to improve training speed with FSDP specifically in the multinode case, and found #1287 which introduced the shard_model utility.

It looks like this removed the ability to specify hybrid sharding, which I believe would be quite useful in the multi-node case (I am seeing very significant TPS decrease in multi-node).

Are there plans to re-introduce hybrid sharding (or is there some other way to accomplish it)? Or is this an area that FSDP2 doesn't have parity yet?

[joecummings]

@ebsmothers I was looking for ways to improve training speed with FSDP specifically in the multinode case, and found #1287 which introduced the shard_model utility.

It looks like this removed the ability to specify hybrid sharding, which I believe would be quite useful in the multi-node case (I am seeing very significant TPS decrease in multi-node).

Are there plans to re-introduce hybrid sharding (or is there some other way to accomplish it)? Or is this an area that FSDP2 doesn't have parity yet?

To my knowledge, there is no specific plan to re-add an option for specifically HYBRID_SHARD. However, if you look at the documented differences between FSDP1 and FSDP2, you'll see that using a 2D mesh and toggling reshard_after_forward is equivalent to using HYBRID_SHARD w/ MiCS or ZERO2 respectively: https://github.com/pytorch/torchtitan/blob/62743776a80afa14975026060a96cf74e40550b8/docs/fsdp.md#fsdp1--fsdp2-api-differences.

What are some of the numbers you're seeing between single node and multi-node TPS?

[nathan-az]

using a 2D mesh and toggling reshard_after_forward is equivalent to using HYBRID_SHARD w/ MiCS or ZERO2 respectively

I assume 2D mesh refers to having a dp_replicate dimension - my question (poorly phrased) was whether there are plans to support this in torchtune? I'm trying to hack this in at the moment by porting from torchtitan, but it would take me a fair bit of time.

What are some of the numbers you're seeing between single node and multi-node TPS?

I'll follow up on this when I get a chance with some figures 🙂 Unfortunately I don't manage the cluster I'm using, so it's difficult to isolate how much is an implementation problem and how much is an inter-node bandwidth option in the cluster setup. I don't believe my cluster has InfiniBand, for example.

[pbontrager]

I think we'd want to benchmark the difference between only FSDP with hybrid (per-node) sharding on a cluster vs 2D FSDP + TP across a cluster. I imagine that for small enough models the only FSDP solution would be faster so it's worth investigating if this makes sense to support better.

[dz1iang]

using a 2D mesh and toggling reshard_after_forward is equivalent to using HYBRID_SHARD w/ MiCS or ZERO2 respectively

I assume 2D mesh refers to having a dp_replicate dimension - my question (poorly phrased) was whether there are plans to support this in torchtune? I'm trying to hack this in at the moment by porting from torchtitan, but it would take me a fair bit of time.

What are some of the numbers you're seeing between single node and multi-node TPS?

I'll follow up on this when I get a chance with some figures 🙂 Unfortunately I don't manage the cluster I'm using, so it's difficult to isolate how much is an implementation problem and how much is an inter-node bandwidth option in the cluster setup. I don't believe my cluster has InfiniBand, for example.

@nathan-az I have plans to conduct multi-machine training, and I have requirements for the training speed. I can assist you with the testing. If it's convenient for you.

[nathan-az]

Thanks @dz1iang! The main issue I've found in my testing (I think correctly using the sharding method from torchtitan) is that the backward pass runtime appears to scale linearly with my number of nodes. I've made a PR to share, but currently TPS is not scaling well with my implementation.

I thought that the 2D mesh was that model and gradient states would reduce-scatter within their replica after backwards passes, then all-reduce across all DP devices during the optimiser step, reducing inter-node communication (i.e. amortizing the cost of the inter-node communication). But I don't think this is happening, or there's an additional operation somewhere.

@ebsmothers Do you know if the way I'm describing is possible? i.e. if we can use FSDP2's APIs to reduce-scatter within replicas on backward steps, and defer the all-reduce across all DP ranks, deferred until optimiser step?

[ebsmothers]

Hi @nathan-az thanks for flagging this. Let me look into your HSDP implementation. If I can make one more request, can you also run with gradient accumulation disabled? I notice some periodicity in your bwd_time plot and we do not currently run our non-stepping backward passes under a no_sync context (as is done in other places). Also I believe titan doesn't do gradient accumulation so turning this off would help ensure we can do a proper apples-to-apples comparison with their version.

Also cc distributed expert @tianyu-l to confirm whether the implementation in your PR #2415 looks correct.

[awgu]

The straightforward comparison is that there is currently no way to support ZeRO-1.

@dz1iang ZeRO-1 is the same thing as ZeRO-2 when there is no microbatching. If there is microbatching, then you can still do ZeRO-1 with FSDP2 via:

is_last_microbatch = (microbatch_idx == num_microbatches - 1)
assert isinstance(model, FSDPModule)
model.set_requires_gradient_sync(is_last_microbatch)
model.set_reshard_after_backward(is_last_microbatch)

(see this unit test as an example)

If you set_requires_gradient_sync(False), then FSDP2 will skip reduce-scatter. If you set_reshard_after_backward(False), then FSDP2 will skip resharding the parameters in backward, so it will not need to re-all-gather the parameters in the next microbatch forward.

[nathan-az]

can you also run with gradient accumulation disabled?

Looks relatively stable to me with gradient_accumulation_steps=1.

I also updated my PR to alternate requires_all_reduce and is_last_backward - keeping the reduce-scatter but deferring the all-reduce. I think this way it corresponds to ZeRO-2 in the accumulation case.

Loss curves look identical to without, although someone more familiar should certainly confirm the implementation 😅.