Enable static type checking with Pyrefly #2136

rchen152 · 2025-12-10T09:31:57Z

Enables static type checking of torchtitan with pyrefly. Type checking the code helps catch bugs earlier in the development cycle.

Adds pyrefly to CI, as part of the linting workflow.
Addresses ~100 type errors that can be fixed via local code changes and updates to type annotations, and silences the rest with # pyrefly: ignore suppression comments. Note that 325efd9 contains all of the non-comment changes.

tianyu-l

Thanks, had some comments.

.github/scripts/setup_pyrefly.sh

tianyu-l · 2025-12-10T10:09:43Z

.github/scripts/setup_pyrefly.sh

+uv pip install \
+  torchft-nightly \
+  -r .ci/docker/requirements.txt \
+  -r .ci/docker/requirements-dev.txt \
+  -r .ci/docker/requirements-flux.txt \
+  -r .ci/docker/requirements-transformers-modeling-backend.txt \
+  -r .ci/docker/requirements-vlm.txt \


We probably shouldn't mix these repo dependency install with type checking install. We should maybe put pyrefly install into .requirements-dev.txt

Done. FYI the reason I originally made this a stand-alone script was so that I could install nightly torch and torchao from https://download.pytorch.org/whl/nightly/cu128 in a separate step. With all the dependencies in requirements-dev.txt, we're now installing those deps from pypi, which changes the type checking results a bit. I had to add some new suppressions.

I ended up moving torch and torchao out again to a separate install step - all the other deps are in requirements-dev.txt

.github/workflows/lint.yaml

tianyu-l · 2025-12-10T10:11:46Z

CONTRIBUTING.md

-6. If you haven't already, complete the Contributor License Agreement ("CLA").
+6. Make sure your code type checks:
+   ```
+   .github/scripts/setup_pyrefly.sh  # one-time setup


similarly, should follow pre-commit style

Done, pyrefly is now run as a pre-commit hook.

Just FYI, the pyrefly hook works a little differently from the others - you have to install all the dependencies yourself (including pyrefly) before running it. This is because a type checker needs to be able to see a project's dependencies, so we use language: system to use the local environment rather than an isolated environment managed by pre-commit.

fegin

This is very helpful. This pyrefly PR already catches at least one bug. We should definitely enforce stricter type checking.

scripts/checkpoint_conversion/convert_from_hf.py

fegin · 2025-12-10T17:55:51Z

scripts/checkpoint_conversion/convert_to_hf.py


    with torch.device("cpu"):
        model = train_spec.model_cls(model_args)
+    # pyrefly: ignore [bad-argument-type]


TODO for me: Change ModelProtocol to include nn.Module as the base class.

@fegin Would you like me to put these TODOs in the code?

fegin · 2025-12-10T17:56:40Z

scripts/checkpoint_conversion/numerical_tests_example.py


    # Calculate KL Divergence
-    kl_loss = F.kl_div(probs1, probs2, "mean")
+    kl_loss = F.kl_div(probs1, probs2, reduction="mean")


This is interesting, so what's the error pyrefly reports without this change?

This is the error:

ERROR Argument `Literal['mean']` is not assignable to parameter `size_average` with type `bool | None` in function `torch.nn.functional.kl_div` [bad-argument-type] --> scripts/checkpoint_conversion/numerical_tests_example.py:28:40 | 28 | kl_loss = F.kl_div(probs1, probs2, "mean") | ^^^^^^ |

fegin · 2025-12-10T17:59:28Z

scripts/estimate/estimation.py

+            optimizers.optimizers,
+            job_config.lr_scheduler,


TODO for me: build_lr_schedulers has an incorrect typing.

torchtitan/components/ft/manager.py

fegin · 2025-12-10T18:02:53Z

torchtitan/components/checkpoint.py

                state_dict,
                storage_writer=storage_writer,
                checkpoint_id=checkpoint_save_id,
+                # pyrefly: ignore [bad-argument-type]


TODO for me: why self.pg has incorrect typing or why async_save doesn't accept the PG typing.

fegin · 2025-12-10T18:05:28Z

torchtitan/components/dataloader.py

        ...


+# pyrefly: ignore [inconsistent-inheritance]


TODO for me: this seems to be an inconsistency issue between checkpoint Stateful and dataloader StatefulDataLoader, which we may not be able to change.

fegin · 2025-12-10T18:10:01Z

torchtitan/components/optimizer.py

 from typing import Any, Generic, Iterator, TypeVar

 import torch
+import torch.distributed.tensor


TODO for me: change the reference of orch.distributed.tensor.DTensor to DTensor.

fegin · 2025-12-10T18:20:05Z

torchtitan/distributed/pipeline_parallel.py


+    # pyrefly: ignore [bad-instantiation]
    schedule = schedule_class(
+        # pyrefly: ignore [bad-argument-type]


typing issue of get_schedule_class, cc., @wconstab @H-Huang

fegin · 2025-12-10T18:31:03Z

torchtitan/distributed/utils.py

            "Context Parallel API. Please update to a newer version."
-        )
+        ) from e



TODO for me: set_rotate_method has been in the main PyTorch for a while. We should be able to remove this try/except.

wwwjn · 2025-12-10T20:03:43Z

.ci/docker/requirements-dev.txt

 pre-commit
+pyrefly==0.45.1
 tomli-w >= 1.1.0
+torch


We can not install pytorch using pip install, as torchtitan depends on pytorch-nightly: https://github.com/pytorch/torchtitan/blob/refs/heads/main/README.md#nightly-builds

@wwwjn I originally had a separate setup_pyrefly.sh script that installed nightly torch and torchao, which I removed based on @tianyu-l's comment that it would be preferable to have all of the dependencies pyrefly needs in requirements-dev.txt. Do you have any suggestions for how to resolve this?

Thank you! I think other dependencies are fine, but for torch and torchao, can we move these 2 package to lint.yaml, and use commands like pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126 --force-reinstall to install them separately?

Done, thanks!

wwwjn · 2025-12-10T20:05:11Z

scripts/checkpoint_conversion/convert_from_hf.py

+    # pyrefly: ignore [bad-argument-type]
    model = ModelWrapper(model)

+    # pyrefly: ignore [not-callable]


pyrefly: ignore [bad-argument-type]

I feel like we need to fix/refactor the code if it doesn't pass pyrefly check, instead of leaving a lot of comments to suppress the error

@wwwjn Adding ignores to get the code to a clean state and then incrementally fixing them is the standard way to enable a type checker on a large code base. I already made fixes to remove 100+ pyrefly: ignores, but the remaining issues look fairly tricky to resolve and IMO would be better tackled as follow-ups rather than trying to combine larger and more risky refactors with enabling the type checker.

rchen152 · 2025-12-10T20:35:40Z

Thanks for all the comments! I believe this is ready for another round of review.

.ci/docker/requirements-dev.txt

tianyu-l · 2025-12-11T10:24:16Z

.ci/docker/requirements-dev.txt

+pyrefly==0.45.1
 tomli-w >= 1.1.0
+torchdata >= 0.8.0
+torchft-nightly


for now we can do similar things to torchao
@fegin we should move torchtft to experiments

tianyu-l · 2025-12-11T10:25:42Z

.github/workflows/lint.yaml

        run: |
-          python -m pip install pre-commit
+          python -m pip install -r requirements-dev.txt
+          python -m pip install --force-reinstall --pre --index-url https://download.pytorch.org/whl/nightly/cu126 torch torchao


Can user not run pyrefly / pre-commit without installing the optional torchao package?

I updated the pyrefly config so that torchao is no longer required - pyrefly now uses ignore-missing-imports for torchao, so it'll replace the import with Any rather than erroring when it's not present.

CONTRIBUTING.md

scripts/checkpoint_conversion/convert_from_hf.py

rchen152 added 4 commits December 5, 2025 13:12

Enable type checking.

87dbb9e

Merge branch 'main' into pyrefly

fd5320c

Update pyrefly version and suppressions.

647b0d8

Fixes

325efd9

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Dec 10, 2025

rchen152 marked this pull request as ready for review December 10, 2025 10:07

rchen152 requested review from fegin, tianyu-l, wconstab and wwwjn as code owners December 10, 2025 10:07

tianyu-l requested changes Dec 10, 2025

View reviewed changes

fegin reviewed Dec 10, 2025

View reviewed changes

Move pyrefly deps into requirements-dev.txt, add pre-commit hook.

f724ebd

wwwjn reviewed Dec 10, 2025

View reviewed changes

rchen152 added 2 commits December 10, 2025 12:20

Fix imports.

16cb3dc

Fix install of torch and torchao.

d0a1e2f

rchen152 requested a review from tianyu-l December 10, 2025 20:35

Remove one more unnecessary import.

cd9484b

tianyu-l reviewed Dec 11, 2025

View reviewed changes

rchen152 added 8 commits December 11, 2025 17:56

Merge branch 'main' into pyrefly

e83adcf

Remove requirements.txt deps from requirements-dev.txt.

003704a

Remove import.

d27eba4

Stop checking torchao dep.

8fa1bf9

Add back missing dep.

1762763

Try removing torchft-nightly dependency.

acce4dd

Fix formatting

22c9682

Fix circular import.

50cd8f6

Enable static type checking with Pyrefly #2136

Are you sure you want to change the base?

Enable static type checking with Pyrefly #2136

Conversation

rchen152 commented Dec 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fegin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wwwjn Dec 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rchen152 commented Dec 10, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

rchen152 commented Dec 10, 2025 •

edited

Loading

wwwjn Dec 10, 2025 •

edited

Loading