[v0.5.10][4] Fix _no_split_modules set not subscriptable in transformers >=5.0

[yueming-yuan]

Summary

• model._no_split_modules changed from list to set in transformers >=5.0
• set doesn't support [0] indexing, causing TypeError: 'set' object is not subscriptable
• All FSDP CI tests fail at model initialization

Test plan

• FSDP CI tests pass

[gemini-code-assist]

Code Review

This pull request updates the apply_fsdp2 function to handle _no_split_modules as either a list or a set, ensuring compatibility with different versions of the transformers library. The reviewer suggested explicitly converting layer_cls_to_wrap to a set to improve the efficiency of membership checks in subsequent operations.

[gemini-code-assist]

While next(iter(layer_cls_to_wrap)) correctly handles the compatibility between list and set types for _no_split_modules in transformers >= 5.0, it is more idiomatic and efficient to ensure layer_cls_to_wrap is a set once at the beginning. This would also optimize the membership check in the subsequent loop (line 689) for cases where transformers < 5.0 provides a list (making it $O(1)$ instead of $O(N)$).

Suggested change

	assert len(layer_cls_to_wrap) > 0 and next(iter(layer_cls_to_wrap)) is not None
	layer_cls_to_wrap = set(model._no_split_modules)
	assert layer_cls_to_wrap and next(iter(layer_cls_to_wrap)) is not None