first commit

tests/trainer/distributed/test_trainer_distributed_fsdp.py

@@ -0,0 +1,668 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+FSDP-specific distributed trainer tests.
+"""
+
+import itertools
+import json
+import os
+import unittest
+from functools import partial
+from pathlib import Path
+from unittest.mock import patch
+
+from parameterized import parameterized
+
+from tests.trainer.trainer_test_utils import TrainerIntegrationCommon, get_regression_trainer  # noqa
+from transformers import HfArgumentParser, PreTrainedConfig, TrainingArguments, is_torch_available
+from transformers.testing_utils import (
+    TestCasePlus,
+    backend_device_count,
+    execute_subprocess_async,
+    get_torch_dist_unique_port,
+    mockenv_context,
+    require_torch,
+    require_torch_accelerator,
+    require_torch_multi_accelerator,
+    slow,
+    torch_device,
+)
+from transformers.trainer_utils import set_seed
+from transformers.utils import (
+    is_torch_bf16_available_on_device,
+    is_torch_fp16_available_on_device,
+)
+
+from .test_trainer_distributed import CONFIGS_DIR, SCRIPTS_DIR, TRAIN_SCRIPT, TrainerDistributedCommon
+
+
+if is_torch_available():
+    import torch
+    from torch import nn
+
+    from transformers import PreTrainedModel
+    from transformers.trainer import FSDP_MODEL_NAME
+
+# Base accelerate configs (version only — model-specific settings via launch args)
+FSDP_CONFIG_FILE = os.path.join(CONFIGS_DIR, "fsdp.yaml")
+FSDP2_CONFIG_FILE = os.path.join(CONFIGS_DIR, "fsdp2.yaml")
+FSDP2_CP_CONFIG_FILE = os.path.join(CONFIGS_DIR, "fsdp2_cp.yaml")
+FSDP_GENERATE_SCRIPT = os.path.join(SCRIPTS_DIR, "fsdp_generate.py")
+
+FSDP_CONFIGS = {
+    "fsdp1": FSDP_CONFIG_FILE,
+    "fsdp2": FSDP2_CONFIG_FILE,
+}
+
+# Launch args shared by all training tests
+TRAIN_LAUNCH_ARGS = [
+    "--fsdp_auto_wrap_policy",
+    "TRANSFORMER_BASED_WRAP",
+]
+
+dtypes = []
+if is_torch_bf16_available_on_device(torch_device):
+    dtypes += ["bf16"]
+if is_torch_fp16_available_on_device(torch_device):
+    dtypes += ["fp16"]
+
+sharding_strategies = ["full_shard", "shard_grad_op"]  # zero3 and zero2
+fsdp_versions = ["fsdp1", "fsdp2"]
+
+config_params = list(itertools.product(sharding_strategies, dtypes))
+# Mixed precision: model loaded in fp32, training with --bf16/--fp16
+mixed_precision_params = list(itertools.product(sharding_strategies, dtypes, fsdp_versions))
+# Pure dtype: model loaded in target dtype, no mixed precision flags
+pure_dtype_params = list(itertools.product(["fp32"] + dtypes, fsdp_versions))
+
+resume_params = [
+    ("FULL_STATE_DICT", "fsdp1"),  # FULL_STATE_DICT only supported for fsdp1
+    ("SHARDED_STATE_DICT", "fsdp1"),
+    ("SHARDED_STATE_DICT", "fsdp2"),
+]
+
+set_seed(42)
+
+
+if is_torch_available():
+    # hack to restore original logging level pre #21700
+    get_regression_trainer = partial(get_regression_trainer, log_level="info")
+
+
+if is_torch_available():
+
+    class _BaseModel(PreTrainedModel):
+        base_model_prefix = "base"
+        config_class = PreTrainedConfig
+
+        def __init__(self, config):
+            super().__init__(config)
+            self.linear = nn.Linear(5, 5)
+            self.linear_2 = nn.Linear(5, 5)
+            self.post_init()
+
+        def forward(self, x):
+            return self.linear_2(self.linear(x))
+
+
+@require_torch
+class InitializeMissingKeysTest(unittest.TestCase):
+    """Tests for FSDP non-rank-0 weight initialization: params should be moved from meta to CPU
+    and marked as initialized without being re-initialized."""
+
+    def _clear_init_flags(self, model):
+        for module in model.modules():
+            if hasattr(module, "_is_hf_initialized"):
+                delattr(module, "_is_hf_initialized")
+        for param in model.parameters():
+            if hasattr(param, "_is_hf_initialized"):
+                delattr(param, "_is_hf_initialized")
+        for buffer in model.buffers():
+            if hasattr(buffer, "_is_hf_initialized"):
+                delattr(buffer, "_is_hf_initialized")
+
+    def test_move_missing_keys_fsdp_non_rank0_moves_meta_to_cpu(self):
+        """FSDP non-rank-0 path should move all params from meta to CPU."""
+        with torch.device("meta"):
+            model = _BaseModel(PreTrainedConfig())
+
+        for param in model.parameters():
+            self.assertEqual(param.device, torch.device("meta"))
+
+        with (
+            patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True),
+            patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False),
+        ):
+            model._move_missing_keys_from_meta_to_device(
+                missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None
+            )
+
+        for name, param in model.named_parameters():
+            self.assertEqual(param.device, torch.device("cpu"), f"param {name} should be on CPU after FSDP move")
+
+    def test_fsdp_non_rank0_end_to_end_no_reinit(self):
+        """End-to-end: move from meta + _initialize_missing_keys should mark all params initialized
+        without changing their values."""
+        with torch.device("meta"):
+            model = _BaseModel(PreTrainedConfig())
+
+        with (
+            patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True),
+            patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False),
+        ):
+            model._move_missing_keys_from_meta_to_device(
+                missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None
+            )
+            pre_init_values = {name: param.clone() for name, param in model.named_parameters()}
+            self._clear_init_flags(model)
+            model._initialize_missing_keys(is_quantized=False)
+
+        for name, param in model.named_parameters():
+            self.assertTrue(getattr(param, "_is_hf_initialized", False), f"param {name} not marked initialized")
+            torch.testing.assert_close(param, pre_init_values[name], msg=f"param {name} was re-initialized")
+        self.assertTrue(getattr(model, "_is_hf_initialized", False))
+
+
+def _parameterized_custom_name_func(func, param_num, param):
+    # customize the test name generator function as we want both params to appear in the sub-test
+    # name, as by default it shows only the first param
+    param_based_name = parameterized.to_safe_name("_".join(str(x) for x in param.args))
+    return f"{func.__name__}_{param_based_name}"
+
+
+# ---------------------------------------------------------------------------
+# Command mixins
+# ---------------------------------------------------------------------------
+
+
+class FSDPCommandsMixin:
+    """Provides ``get_torchrun_cmd`` and ``get_accelerate_cmd`` for FSDP."""
+
+    def get_torchrun_cmd(self, script, script_args=None, num_processes=None):
+        if num_processes is None:
+            num_processes = backend_device_count(torch_device)
+        port = get_torch_dist_unique_port()
+        cmd = [
+            "torchrun",
+            f"--nproc_per_node={num_processes}",
+            "--nnodes=1",
+            f"--master_port={port}",
+            script,
+        ]
+        if script_args:
+            cmd.extend(script_args)
+        return cmd
+
+    def get_accelerate_cmd(
+        self, script, config_file, launch_args=None, script_args=None, num_processes=None, **kwargs
+    ):
+        if num_processes is None:
+            num_processes = backend_device_count(torch_device)
+        port = get_torch_dist_unique_port()
+        cmd = [
+            "accelerate",
+            "launch",
+            "--config_file",
+            config_file,
+            "--num_processes",
+            str(num_processes),
+            "--main_process_port",
+            str(port),
+        ]
+        if launch_args:
+            cmd.extend(launch_args)
+        cmd.append(script)
+        if script_args:
+            cmd.extend(script_args)
+        return cmd
+
+
+# ---------------------------------------------------------------------------
+# Config parsing tests (no distributed training runs)
+# ---------------------------------------------------------------------------
+
+
+@require_torch_accelerator
+class TestFSDPConfig(TestCasePlus):
+    def setUp(self):
+        super().setUp()
+        master_port = get_torch_dist_unique_port()
+        self.dist_env_1_gpu = {
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(master_port),
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+        }
+        self.accelerate_fsdp_config = {
+            "fsdp_activation_checkpointing": False,
+            "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
+            "fsdp_backward_prefetch": "BACKWARD_PRE",
+            "fsdp_cpu_ram_efficient_loading": True,
+            "fsdp_forward_prefetch": False,
+            "fsdp_offload_params": False,
+            "fsdp_reshard_after_forward": "FULL_SHARD",
+            "fsdp_state_dict_type": "FULL_STATE_DICT",
+            "fsdp_sync_module_states": True,
+            "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
+            "fsdp_use_orig_params": True,
+            "fsdp_version": 1,
+        }
+
+        self.fsdp_config = {
+            "backward_prefetch": "BACKWARD_PRE",
+            "forward_prefetch": "false",
+            "limit_all_gathers": "false",
+            "use_orig_params": "true",
+            "sync_module_states": "true",
+            "cpu_ram_efficient_loading": "true",
+            "activation_checkpointing": "false",
+            "min_num_params": 1,
+        }
+
+    @parameterized.expand(config_params, name_func=_parameterized_custom_name_func)
+    def test_accelerate_fsdp_config(self, sharding_strategy, dtype):
+        output_dir = self.get_auto_remove_tmp_dir()
+        # Snapshot before trainer construction — `_process_fsdp_args` strips the
+        # `fsdp_` prefix in place.
+        expected = dict(self.accelerate_fsdp_config)
+        kwargs = {
+            "output_dir": output_dir,
+            "train_len": 128,
+            "save_steps": 5,
+            "learning_rate": 0.1,
+            "fsdp": f"{sharding_strategy} offload auto_wrap",
+            "fsdp_config": self.accelerate_fsdp_config,
+        }
+        kwargs[dtype] = True
+        with mockenv_context(**self.dist_env_1_gpu):
+            trainer = get_regression_trainer(**kwargs)
+            self.assertIs(trainer.args.fsdp, True)
+            self.assertTrue(trainer.args.fsdp_config.get("cpu_offload"))
+            for k, v in expected.items():
+                assert k.startswith("fsdp_")
+                # `transformer_layer_cls_to_wrap` is normalized from str → list during parsing.
+                if k == "fsdp_transformer_layer_cls_to_wrap" and isinstance(v, str):
+                    v = [v]
+                self.assertEqual(trainer.args.fsdp_config[k[5:]], v)
+
+    def test_torchrun_fsdp_config(self):
+        """Verify that --fsdp + --fsdp_config (torchrun-style) are parsed correctly."""
+        output_dir = self.get_auto_remove_tmp_dir()
+        fsdp_config = {"fsdp_transformer_layer_cls_to_wrap": "Qwen2DecoderLayer"}
+        kwargs = {
+            "output_dir": output_dir,
+            "train_len": 128,
+            "save_steps": 5,
+            "learning_rate": 0.1,
+            "fsdp": "full_shard auto_wrap",
+            "fsdp_config": fsdp_config,
+            "bf16": True,
+        }
+        with mockenv_context(**self.dist_env_1_gpu):
+            trainer = get_regression_trainer(**kwargs)
+            self.assertIs(trainer.args.fsdp, True)
+            # fsdp_ prefix is stripped and value is normalized to a list during parsing
+            self.assertIn("Qwen2DecoderLayer", trainer.args.fsdp_config["transformer_layer_cls_to_wrap"])
+
+    def test_fsdp_cli_parsing(self):
+        """`--fsdp` (bare) → True; legacy `--fsdp full_shard` still parses; absent → None."""
+        parser = HfArgumentParser(TrainingArguments)
+        base = ["--output_dir", "/tmp/x"]
+
+        args, _ = parser.parse_known_args([*base, "--fsdp"])
+        self.assertIs(args.fsdp, True)
+
+        args, _ = parser.parse_known_args([*base, "--fsdp", "full_shard"])
+        self.assertEqual(args.fsdp, "full_shard")
+
+        args, _ = parser.parse_known_args(base)
+        self.assertIsNone(args.fsdp)
+
+        # Bare `--fsdp` should resolve to a fully enabled FSDP setup through `_process_fsdp_args`.
+        with mockenv_context(**self.dist_env_1_gpu):
+            trainer_args = TrainingArguments(output_dir="/tmp/x", fsdp=True)
+            self.assertIs(trainer_args.fsdp, True)
+            self.assertIsNotNone(trainer_args.fsdp_plugin_args)
+
+    @parameterized.expand(config_params, name_func=_parameterized_custom_name_func)
+    def test_fsdp_config(self, sharding_strategy, dtype):
+        output_dir = self.get_auto_remove_tmp_dir()
+        kwargs = {
+            "output_dir": output_dir,
+            "train_len": 128,
+            "save_steps": 5,
+            "learning_rate": 0.1,
+            "fsdp": f"{sharding_strategy} offload auto_wrap",
+            "fsdp_config": self.fsdp_config,
+        }
+        kwargs[dtype] = True
+        with mockenv_context(**self.dist_env_1_gpu):
+            trainer = get_regression_trainer(**kwargs)
+            self.assertIs(trainer.args.fsdp, True)
+            self.assertTrue(trainer.args.fsdp_config.get("cpu_offload"))
+            for k, v in self.fsdp_config.items():
+                self.assertEqual(trainer.args.fsdp_config[k], v)
+
+
+# ---------------------------------------------------------------------------
+# FSDP distributed tests
+# ---------------------------------------------------------------------------
+
+
+@require_torch_multi_accelerator
+class TestTrainerDistributedFSDP(FSDPCommandsMixin, TestCasePlus):
+    def _run_env_check(self, cmd, num_processes):
+        """Run the env check script and return per-rank results."""
+        execute_subprocess_async(cmd, env=self.get_env())
+        # output_dir is always the last script_arg value
+        output_dir = cmd[cmd.index("--output_dir") + 1]
+        results = []
+        for rank in range(num_processes):
+            with open(os.path.join(output_dir, f"env_rank{rank}.json")) as f:
+                results.append(json.load(f))
+        return results
+
+    def test_torchrun_accelerate_fsdp1_env_parity(self):
+        """Verify torchrun+--fsdp and accelerate launch produce the same FSDP1 env."""
+        script = os.path.join(SCRIPTS_DIR, "torchrun_env_check.py")
+        num_processes = backend_device_count(torch_device)
+
+        torchrun_dir = self.get_auto_remove_tmp_dir()
+        torchrun_results = self._run_env_check(
+            self.get_torchrun_cmd(
+                script,
+                script_args=[
+                    "--output_dir",
+                    torchrun_dir,
+                    "--fsdp",
+                    "full_shard",
+                    "--fsdp_config",
+                    '{"fsdp_version": 1}',
+                ],
+                num_processes=num_processes,
+            ),
+            num_processes,
+        )
+
+        accel_dir = self.get_auto_remove_tmp_dir()
+        accel_results = self._run_env_check(
+            self.get_accelerate_cmd(
+                script, FSDP_CONFIG_FILE, script_args=["--output_dir", accel_dir], num_processes=num_processes
+            ),
+            num_processes,
+        )
+
+        self._check_parity(torchrun_results, accel_results, num_processes, expected_fsdp_version=1)
+
+    def test_torchrun_accelerate_fsdp2_env_parity(self):
+        """Verify torchrun+--fsdp and accelerate launch produce the same FSDP2 env."""
+        script = os.path.join(SCRIPTS_DIR, "torchrun_env_check.py")
+        num_processes = backend_device_count(torch_device)
+
+        torchrun_dir = self.get_auto_remove_tmp_dir()
+        torchrun_results = self._run_env_check(
+            self.get_torchrun_cmd(
+                script,
+                script_args=[
+                    "--output_dir",
+                    torchrun_dir,
+                    "--fsdp",
+                    "full_shard",
+                    "--fsdp_config",
+                    '{"fsdp_version": 2}',
+                ],
+                num_processes=num_processes,
+            ),
+            num_processes,
+        )
+
+        accel_dir = self.get_auto_remove_tmp_dir()
+        accel_results = self._run_env_check(
+            self.get_accelerate_cmd(
+                script, FSDP2_CONFIG_FILE, script_args=["--output_dir", accel_dir], num_processes=num_processes
+            ),
+            num_processes,
+        )
+
+        self._check_parity(torchrun_results, accel_results, num_processes, expected_fsdp_version=2)
+
+    def _check_parity(self, torchrun_results, accel_results, num_processes, expected_fsdp_version):
+        for rank in range(num_processes):
+            tr, ac = torchrun_results[rank], accel_results[rank]
+
+            # Both should agree on distributed env
+            self.assertEqual(tr["args_world_size"], ac["args_world_size"])
+            self.assertEqual(tr["args_process_index"], ac["args_process_index"])
+            self.assertEqual(tr["args_parallel_mode"], ac["args_parallel_mode"])
+            self.assertEqual(tr["accelerator_num_processes"], ac["accelerator_num_processes"])
+            self.assertEqual(tr["accelerator_use_distributed"], ac["accelerator_use_distributed"])
+
+            for info in (tr, ac):
+                # Rank consistency across all layers
+                self.assertEqual(info["env_world_size"], str(num_processes))
+                self.assertEqual(info["env_rank"], str(rank))
+                self.assertEqual(info["args_process_index"], rank)
+                self.assertEqual(info["args_local_process_index"], rank)
+                self.assertEqual(info["accelerator_process_index"], rank)
+                self.assertEqual(info["accelerator_local_process_index"], rank)
+                self.assertEqual(info["args_n_gpu"], 1)
+                self.assertEqual(info["accelerator_is_main_process"], rank == 0)
+                self.assertEqual(info["accelerator_is_local_main_process"], rank == 0)
+                self.assertIn(f"{torch_device}:{rank}", info["accelerator_device"])
+
+                # Both should have FSDP enabled with the correct version
+                self.assertEqual(info["accelerator_distributed_type"], "DistributedType.FSDP")
+                self.assertTrue(info["trainer_is_fsdp_enabled"])
+                self.assertFalse(info["trainer_is_deepspeed_enabled"])
+                self.assertEqual(info["fsdp_version"], expected_fsdp_version)
+                self.assertNotIn("deepspeed_zero_stage", info)
+
+
+# ---------------------------------------------------------------------------
+# All distributed FSDP training tests
+# ---------------------------------------------------------------------------
+@slow
+@require_torch_multi_accelerator
+class TestTrainerDistributedFSDPCommon(
+    FSDPCommandsMixin, TrainerDistributedCommon, TestCasePlus, TrainerIntegrationCommon
+):
+    # -------------------------------------------------------------------
+    # FSDP training — accelerate (parameterized over fsdp version)
+    # -------------------------------------------------------------------
+
+    # Pure dtype training: model loaded in target dtype, no mixed precision
+    @parameterized.expand(pure_dtype_params, name_func=_parameterized_custom_name_func)
+    def test_training(self, dtype, fsdp_version):
+        self.check_training(dtype, config_file=FSDP_CONFIGS[fsdp_version])
+
+    # Mixed precision: model loaded in fp32, training with --bf16/--fp16
+    @parameterized.expand(mixed_precision_params, name_func=_parameterized_custom_name_func)
+    def test_training_mixed_precision(self, sharding_strategy, dtype, fsdp_version):
+        if fsdp_version == "fsdp2":
+            reshard = "true" if sharding_strategy == "full_shard" else "false"
+        else:
+            reshard = sharding_strategy.upper()
+        launch_args = list(TRAIN_LAUNCH_ARGS) + ["--fsdp_reshard_after_forward", reshard]
+        self.check_mixed_precision(dtype, config_file=FSDP_CONFIGS[fsdp_version], launch_args=launch_args)
+
+    @parameterized.expand(["true", "false"], name_func=_parameterized_custom_name_func)
+    def test_fsdp2_cpu_ram_efficient_loading(self, cpu_ram_efficient_loading):
+        launch_args = list(TRAIN_LAUNCH_ARGS) + [
+            "--fsdp_cpu_ram_efficient_loading",
+            cpu_ram_efficient_loading,
+        ]
+        self.check_training("bf16", config_file=FSDP2_CONFIG_FILE, launch_args=launch_args)
+
+    @parameterized.expand(fsdp_versions, name_func=_parameterized_custom_name_func)
+    def test_training_with_gradient_accumulation(self, fsdp_version):
+        self.check_gradient_accumulation(config_file=FSDP_CONFIGS[fsdp_version])
+
+    @parameterized.expand(fsdp_versions, name_func=_parameterized_custom_name_func)
+    def test_basic_run_with_cpu_offload(self, fsdp_version):
+        output_dir = self.get_auto_remove_tmp_dir()
+        args = self._get_default_script_args(output_dir) + ["--bf16", "--max_steps", "10"]
+        launch_args = list(TRAIN_LAUNCH_ARGS) + ["--fsdp_offload_params", "true"]
+        execute_subprocess_async(
+            self.get_accelerate_cmd(
+                TRAIN_SCRIPT, script_args=args, config_file=FSDP_CONFIGS[fsdp_version], launch_args=launch_args
+            ),
+            env=self.get_env(),
+        )
+
+    @parameterized.expand(resume_params, name_func=_parameterized_custom_name_func)
+    def test_training_and_can_resume_normally(self, state_dict_type, fsdp_version):
+        output_dir = self.get_auto_remove_tmp_dir()
+        args = self._get_default_script_args(output_dir, num_epochs=2, logging_steps=2, save_steps=2)
+
+        launch_args = list(TRAIN_LAUNCH_ARGS) + ["--fsdp_state_dict_type", state_dict_type]
+        cmd_kwargs = {"config_file": FSDP_CONFIGS[fsdp_version], "launch_args": launch_args}
+
+        logs = self._train_and_get_log_history(
+            self.get_accelerate_cmd(TRAIN_SCRIPT, script_args=args, **cmd_kwargs),
+            output_dir,
+        )
+
+        # resume from ckpt
+        checkpoint = os.path.join(output_dir, "checkpoint-2")
+        resume_args = args + ["--resume_from_checkpoint", checkpoint]
+
+        is_fsdp_ckpt = os.path.isdir(checkpoint) and (
+            # this checks the FSDP state dict when `SHARDED_STATE_DICT` is used
+            any(
+                FSDP_MODEL_NAME in folder_name
+                for folder_name in os.listdir(checkpoint)
+                if os.path.isdir(os.path.join(checkpoint, folder_name))
+            )
+            # this checks the FSDP state dict when `FULL_STATE_DICT` is used
+            or os.path.isfile(os.path.join(checkpoint, f"{FSDP_MODEL_NAME}.bin"))
+        )
+        self.assertTrue(is_fsdp_ckpt)
+
+        logs_resume = self._train_and_get_log_history(
+            self.get_accelerate_cmd(TRAIN_SCRIPT, script_args=resume_args, **cmd_kwargs),
+            output_dir,
+        )
+
+        for log, log1 in zip(logs, logs_resume):
+            if "learning_rate" in log:
+                self.assertAlmostEqual(log["learning_rate"], log1["learning_rate"], delta=1e-5)
+
+    # -------------------------------------------------------------------
+    # Context parallel tests
+    # -------------------------------------------------------------------
+    def test_cp_equivalence(self):
+        """Test that CP produces the same losses as without CP."""
+
+        # CP doesn't work with Qwen2 (DTensor mixing error), so we use Llama here.
+        launch_args = list(TRAIN_LAUNCH_ARGS) + ["--fsdp_state_dict_type", "SHARDED_STATE_DICT"]
+        cp_script_args = [
+            "--model_name",
+            "hf-internal-testing/tiny-random-LlamaForCausalLM",
+            "--max_steps",
+            "10",
+            "--per_device_train_batch_size",
+            "1",
+            "--seed",
+            "42",
+            "--logging_steps",
+            "1",
+            "--save_strategy",
+            "no",
+            "--model_dtype",
+            "fp32",
+            "--attn_implementation",
+            "sdpa",
+            "--pad_to_multiple_of",
+            "4",
+        ]
+
+        # Step 1: Run with CP enabled (cp_size=2)
+        cp_yes_output_dir = Path(self.get_auto_remove_tmp_dir()).resolve()
+        cp_yes_losses_path = cp_yes_output_dir / "cp_yes_losses.json"
+        cmd = self.get_accelerate_cmd(
+            TRAIN_SCRIPT,
+            config_file=FSDP2_CP_CONFIG_FILE,
+            launch_args=launch_args,
+            script_args=["--output_dir", str(cp_yes_output_dir), "--loss_output_file", str(cp_yes_losses_path)]
+            + cp_script_args,
+        )
+        execute_subprocess_async(cmd, env=self.get_env())
+
+        # Step 2: Run without CP (FSDP with num_processes=1, no parallelism_config)
+        cp_no_output_dir = Path(self.get_auto_remove_tmp_dir()).resolve()
+        cp_no_losses_path = cp_no_output_dir / "cp_no_losses.json"
+
+        cmd = self.get_accelerate_cmd(
+            TRAIN_SCRIPT,
+            config_file=FSDP2_CONFIG_FILE,
+            launch_args=launch_args,
+            script_args=[
+                "--output_dir",
+                str(cp_no_output_dir),
+                "--loss_output_file",
+                str(cp_no_losses_path),
+            ]
+            + cp_script_args,
+            num_processes=1,
+        )
+        execute_subprocess_async(cmd, env=self.get_env())
+
+        # Compare losses
+        with open(cp_yes_losses_path) as f:
+            cp_yes_losses = json.load(f)
+        with open(cp_no_losses_path) as f:
+            cp_no_losses = json.load(f)
+
+        assert len(cp_yes_losses) == len(cp_no_losses), (
+            f"Different number of losses: CP has {len(cp_yes_losses)}, no-CP has {len(cp_no_losses)}"
+        )
+
+        cp_yes_losses_tensor = torch.tensor(cp_yes_losses)
+        cp_no_losses_tensor = torch.tensor(cp_no_losses)
+
+        torch.testing.assert_close(
+            cp_yes_losses_tensor,
+            cp_no_losses_tensor,
+            rtol=2e-2,
+            atol=2e-2,
+            msg=f"CP losses {cp_yes_losses} do not match non-CP losses {cp_no_losses}",
+        )
+
+    # -------------------------------------------------------------------
+    # FSDP eval tests
+    # -------------------------------------------------------------------
+    def test_eval(self):
+        self.check_eval(config_file=FSDP_CONFIG_FILE)
+
+    # -------------------------------------------------------------------
+    # FSDP generation tests (moved from tests/generation/test_fsdp.py)
+    # -------------------------------------------------------------------
+    def test_fsdp_generate(self):
+        cmd = self.get_accelerate_cmd(
+            FSDP_GENERATE_SCRIPT,
+            config_file=FSDP_CONFIG_FILE,
+            script_args=["--fsdp"],
+        )
+        execute_subprocess_async(cmd, env=self.get_env())
+
+    def test_fsdp2_generate(self):
+        cmd = self.get_accelerate_cmd(
+            FSDP_GENERATE_SCRIPT,
+            config_file=FSDP2_CONFIG_FILE,
+            script_args=["--fsdp2"],
+        )
+        execute_subprocess_async(cmd, env=self.get_env())