Conversation
build_trainer() routes accelerator="xla"/"tpu" through an XLAPrecision("bf16-true")
plugin instead of precision=, since XLAStrategy's precision_plugin setter rejects
standard precision strings. RFDETR._resolve_trainer_device_kwargs() gains an explicit
device.type == "xla" -> accelerator="tpu" arm ahead of the auto-detection warning.
Extends the multi_scale/torch.compile log message with XLA recompilation guidance and
adds a static-shape note to TrainConfig.multi_scale. Adds the first @pytest.mark.xla
test (grid_sample gather-path execution, torch_xla-gated) and flips ci-tests-xla.yml's
exit-5 self-pass guard to enforcing now that a real xla-marked test exists.
---
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Not ready to approve
The XLA plugin/precision logic in build_trainer() can be overridden by trainer_kwargs (reintroducing the PTL XLA precision error) and currently forces bf16 even when AMP is disabled.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
This PR advances XLA/TPU training support by aligning PyTorch Lightning trainer construction with XLAStrategy’s precision requirements, improving device→accelerator mapping for "xla", and adding initial XLA-execution coverage plus CI enforcement for XLA-marked tests.
Changes:
- Route
accelerator=("xla"|"tpu")through anXLAPrecision("bf16-true")plugin instead ofprecision=..., and maptorch.device(...).type == "xla"toaccelerator="tpu". - Add XLA runtime test coverage for
_bilinear_grid_sample(gather-path, no CPU fallback) and trainer/device mapping tests. - Update TPU/XLA guidance in logs/config comments and enforce the XLA CI job now that
@pytest.mark.xlatests exist.
File summaries
| File | Description |
|---|---|
src/rfdetr/training/trainer.py |
Adds XLA/TPU precision handling via XLAPrecision plugin instead of precision=. |
src/rfdetr/detr.py |
Maps device.type=="xla" to PTL accelerator="tpu" before auto-detection warnings. |
src/rfdetr/training/module_model.py |
Extends multi-scale/compile log messaging with XLA recompilation guidance. |
src/rfdetr/config.py |
Adds a TPU/XLA static-shape note to TrainConfig.multi_scale. |
tests/utilities/test_tensors.py |
Adds a real torch_xla execution test asserting gather-path runs without CPU fallback ops. |
tests/training/test_detr_shim.py |
Adds unit tests asserting "xla" device inputs map to accelerator="tpu" (and indices to devices=[...]). |
tests/training/test_build_trainer.py |
Adds unit test asserting XLA accelerators use XLAPrecision plugin and omit precision=. |
.github/workflows/ci-tests-xla.yml |
Removes the “exit code 5 self-pass” guard to enforce XLA test execution. |
Review details
- Files reviewed: 8/8 changed files
- Comments generated: 1
- Review effort level: Low
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
Codecov Report❌ Patch coverage is ❌ Your project check has failed because the head coverage (84%) is below the target coverage (95%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #1257 +/- ##
=======================================
- Coverage 84% 84% -0%
=======================================
Files 108 108
Lines 13455 13478 +23
=======================================
+ Hits 11276 11295 +19
- Misses 2179 2183 +4 🚀 New features to boost your workflow:
|
- Preserve caller plugins while appending the required XLA precision plugin - Ignore incompatible precision kwargs and respect disabled AMP on XLA --- Co-authored-by: Codex <codex@openai.com>
Co-authored-by: Codex <codex@openai.com>
build_trainer() routes accelerator="xla"/"tpu" through an XLAPrecision("bf16-true") plugin instead of precision=, since XLAStrategy's precision_plugin setter rejects standard precision strings. RFDETR._resolve_trainer_device_kwargs() gains an explicit device.type == "xla" -> accelerator="tpu" arm ahead of the auto-detection warning. Extends the multi_scale/torch.compile log message with XLA recompilation guidance and adds a static-shape note to TrainConfig.multi_scale. Adds the first @pytest.mark.xla test (grid_sample gather-path execution, torch_xla-gated) and flips ci-tests-xla.yml's exit-5 self-pass guard to enforcing now that a real xla-marked test exists.
part of #1058