Support FSDPv2 compute dtype

[lausannel]

This PR does the following:

1. Supports using torchdistx for model initialization to reduce memory usage
2. Supports different compute dtypes

[JackCaoG]

AFAIK torchdistx is not being maintained, has that changed?

[lausannel]

You're correct. It's worth noting that torchdistx is still in use, and a similar implementation logic is employed in FSDPv1.

xla/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py

Lines 486 to 490 in 72fda76

	_materialize_module(
	module,
	param_init_fn,
	[], # TODO: ignored_params is set to empty now, pass in correct params when this feature is fully enabled
	deferred_init_check_fn=lambda k: not isinstance(k, wrapper_cls))

[JackCaoG]

actually I would prefer not to add this logic here. I think pytorch has been trying to use the meta tensor to solve the parameter issue, I prefer not to add additional logics to a library that has been deprecated.

[lausannel]

I have removed all the torchdistx-related content as per your suggestion. Could you please review the changes again?