fix: add NCCL-safe utilities

[Yunnglin]

Summary

This PR introduces production-grade fault tolerance and observability for distributed training, along with a complete E2E test infrastructure overhaul.

Changes

1. NCCL-safe fault tolerance (src/twinkle/utils/nccl_safe.py)

Three layers of protection to prevent NCCL hangs in production (TWINKLE_FAIL_FAST=0):

• Layer 1 - safe_loss(): Wraps loss instances (as a Loss subclass) to catch computation errors and return graph-connected zero loss for FSDP ReduceScatter compatibility.
• Layer 2 - @nccl_safe: Detects forward/backward boundaries and forces zero-gradient backward when forward ran but backward didn't, preventing DP rank desynchronization.
• Layer 3 - @nccl_safe_megatron: Catches pre-communication errors in Megatron methods where the entire body is NCCL-critical.

2. Deep health check

• Add ping() method to all model backends (mock/transformers/megatron)
• Add /twinkle/healthz/deep gateway endpoint that probes model actor liveness
• Update entrypoint.sh watchdog to check deep health URL
• Returns 503 when model actors are dead (OOM/SIGKILL)

3. DPO ragged logps handling

• Fix forward_backward in Megatron backend: flatten and pad_and_stack ragged ref_outputs['logps'] from variable-length PP microbatches
• Fix _tensor_output_to_rows in common.py: handle ragged list[list] after to_cpu_safe_output serialization

4. Megatron backend sync=True

• Add sync=True to twinkle-native forward_backward @remote_function, consistent with base class and tinker_forward_backward

5. TWINKLE_FAIL_FAST propagation

• Propagate TWINKLE_FAIL_FAST: "0" to all service runtime_env in cookbook configs
• Add env propagation in launcher/env_propagation.py

6. E2E test infrastructure

• New tests/server/integration/e2e_helpers.py with shared utilities
• Backend-agnostic tests: test_sft_e2e.py, test_dpo_e2e.py, test_grpo_e2e.py
• NCCL-safe fault tolerance tests: test_nccl_safe_twinkle_e2e.py, test_nccl_safe_tinker_e2e.py
• Refactored test_full_cycle_e2e.py (checkpoint save/load/resume)
• One-click server startup: tests/server/start_e2e_server.py
• Megatron E2E configs: server_config_4b_e2e_megatron.yaml

Testing

All tests verified on real GPU environment (4x model + 1x sampler):

• SFT E2E (twinkle + tinker, transformers + megatron): loss decreases
• DPO E2E (twinkle + tinker): loss decreases, reward margins increase
• GRPO E2E: sampling + training completes without timeout
• NCCL-safe E2E: graceful degradation under fault injection
• Full cycle: save/load/resume checkpoint correctness

[gemini-code-assist]

Code Review

This pull request introduces NCCL-safe fault tolerance utilities (safe_loss and @nccl_safe) to prevent distributed training hangs by catching errors gracefully and forcing zero-gradient backward passes when necessary. It also propagates the TWINKLE_FAIL_FAST environment variable, updates server configurations, and adds extensive E2E and unit tests. However, several critical issues were identified in the new nccl_safe.py implementation: first, safe_loss returns a function wrapper instead of a Loss instance, which will fail pipeline assertions expecting a Loss subclass; second, the fallback path in _force_zero_backward can raise an AttributeError if model.model is a list, and detaching the parameter breaks the autograd graph connectivity required for FSDP hooks to fire.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Pull request overview

Copilot reviewed 37 out of 37 changed files in this pull request and generated 8 comments.

Comments suppressed due to low confidence (1)

tests/server/integration/test_full_cycle_e2e.py:63

• SAVE_DIR is hard-coded to a /mnt/nas2/... path, which makes this E2E test fail outside that specific environment. Use an env var override and default to a writable temp directory.