def test_tp_train(
    tp_degree: int,
    tp_strategy: str,
    tmp_path: Path,
    tiny_text_dataset_path: Path,
):
    """Test that we can train with FSDP-TP."""
    tp_dataset_name = dist.all_gather_object(tiny_text_dataset_path)[0]

    # Train model with TP and get loss
    tp_cfg = get_cfg(pathlib.Path(tp_dataset_name), tp_strategy, tp_degree)
    tp_trainer = train(tp_cfg)
    tp_trainer.close()
    tp_loss = get_loss_array(tp_trainer)

    # Compare loss and expected loss for TP
    import numpy as np
    expected_tp_loss = np.array([
        11.779396,
        11.750473,
        11.720526,
        11.778467,
        11.735085,
        11.741872,
    ])
    np.testing.assert_allclose(tp_loss, expected_tp_loss)
Test that we can train with FSDP-TP.
tests/tp/test_tp_strategies.py