[SimpleFSDP] Add CI for SimpleFSDP

[ruisizhang123]

This PR adds CI for SimpleFSDP. It mainly adopted FSDP-related integration tests from this file: https://github.com/pytorch/torchtitan/blob/main/tests/integration_tests.py

A few things to list here:

• Currently, DDP/HSDP + TP (SimpleFSDP's replicate and hybrid_shard mode) is not supported. I'll add support for them in follow-up PRs.

• DCP + SimpleFSDP failed in test_generate test on my end. The output seems to indicate some of the model weights are not correctly checkpointed when resharding from [dp:4] -> [tp:4]. [dp:4] -> [dp:2, tp:2] works out.

[rank2]:   File "/home/ruisi/pytorch/torch/distributed/checkpoint/default_planner.py", line 471, in create_default_local_load_plan
[rank2]:     raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
[rank2]: RuntimeError: Missing key in checkpoint state_dict: model.tok_embeddings.weight.

[tianyu-l]

I had some suggestions on truncating the test set: some of them I believe are not proper to run, some else are to relieve CI burden.

For checkpointing, we marked it as WIP feature. I'm a bit surprise it can run -- can we verify the behavior is expected?

[tianyu-l]

This is for testing generate script in https://github.com/pytorch/torchtitan/tree/main/scripts/generate
Let's remove

[ruisizhang123]

I think it's needed to verify if DCP can correctly load checkpoint for this line:

torchtitan/scripts/generate/test_generate.py

Lines 147 to 150 in ed2bbc0

	begin = time.monotonic()
	logger.info(f"Loading chkpt at: {checkpoint_path}")
	dcp.load(state_dict, checkpoint_id=checkpoint_path)
	logger.info(f"Finished loading chkpt in {time.monotonic() - begin:.2f} seconds.")

.

Currently, this line is not correctly loading weights from saved checkpoints. I have commented entry to this test as TODO from L283. Maybe I should comment and add TODO this test as well.

[tianyu-l]

Let's add a test for DCP resharding, similar to https://github.com/pytorch/torchtitan/blob/main/tests/integration_tests.py#L478. If we can reproduce the error, we should remove this test, which is to test the generate script, not (only) for DCP loading.

[ruisizhang123]

I did a quick verification for DCP. It turns out resharding from [dp:4] --> [dp:2, tp:2] still works properly. The loss is converging.

I took a closer look at test_generate. The error comes from resharding [dp:4] --> [tp:4] when initializing the parallel dim here:

torchtitan/scripts/generate/test_generate.py

Lines 121 to 130 in ed2bbc0

	dist_utils.init_distributed(config)
	parallel_dims = ParallelDims(
	dp_replicate=1,
	dp_shard=-1,
	cp=1,
	tp=world_size,
	pp=1,
	world_size=world_size,
	enable_loss_parallel=False,
	)

. I'll add another config under the optional_checkpoint test for this.

[ruisizhang123]

For checkpointing, we marked it as WIP feature. I

For DCP, it will throw error when loading the weight dicts in test_generate test as I mention in the PR description. Some of the weights are not properly checkpointed. We probably need some efforts for the DCP integration.

[fegin]

Does SimpleFSDP change FQNs?

[ruisizhang123]

Does SimpleFSDP change FQNs?

I think SimpleFSDP only parametrizes the parameters but doesn't explicitly change FQNs. (@tianyu-l can confirm this)

However, I found the model weight key is renamed to model.tok_embeddings.parametrizations.weight.original after being wrapped by SimpleFSDP. When the weight is loaded by [tp:4], SimpleFSDP is not called here and the weight is the normal model.tok_embeddings.weight.

This should be the reason causing [dp:4] -> [tp:4] resharding error.

[tianyu-l]

LGTM!

However, I found the model weight key is renamed to model.tok_embeddings.parametrizations.weight.original after being wrapped by SimpleFSDP. When the weight is loaded by [tp:4], SimpleFSDP is not called here and the weight is the normal model.tok_embeddings.weight.

Yes, general DCP compatibility is broken by parametrization. IIRC @fmassa has a fix/improvement to parametrization so it does not interfere with parameter names for DCP. Let's follow up on the fix later.