Guard the cross-attention FSDP2 wrap for pipeline stage modules

[Naeemkh]

Summary

apply_fsdp2 receives a PipelineStageModule under pipeline parallelism, but _fsdp_wrap_transformer_blocks read transformer.cross_attention_layers unconditionally. Only Transformer defines that attribute, so every pp > 1 run died with AttributeError before the first step.

• Guard the cross-attention FSDP2 wrap with getattr. apply_fsdp2 receives a PipelineStageModule under PP, which has no cross_attention_layers, so every pp > 1 run died with AttributeError before the first step.
• Widen both signatures to Transformer | PipelineStageModule. With the annotation widened and the guard removed, pyright flags the offending line, so CI catches the next one.
• Add the first test that calls apply_fsdp2 on a stage module. CPU only, runs in CI today, fails on the unpatched tree.

Testing

• uv run ruff check kempnerforge/ tests/ passes
• uv run ruff format --check kempnerforge/ tests/ scripts/ passes
• uv run pyright kempnerforge/ passes (0 errors)
• uv run pytest tests/unit/ -v --timeout=60 passes
• If distributed code changed: uv run torchrun --nproc_per_node=4 -m pytest tests/distributed/ -v
• If training loop / parallelism / optimizers changed: uv run pytest tests/e2e/ --e2e -v

Closes #185

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/distributed/parallel.py	65.69% <100.00%> (+6.63%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[amazloumi]

LGTM. none of comments are blockers.

[amazloumi]

Thanks for the changes.