Fix NameError in torchtune rlhf/loss/dpo.py from missing imports - #2966
Fix NameError in torchtune rlhf/loss/dpo.py from missing imports#2966wjh70301-meta wants to merge 1 commit into
Conversation
Summary:
`torchtune/rlhf/loss/dpo.py` referenced `dataclass` and `TypeVar` (line `T = TypeVar("T", bound=dataclass)`) without importing them, raising `NameError: name 'TypeVar' is not defined` at import time. This blocked every Mitra DPO unit that imports `torchtune.rlhf.loss.DPOLoss` — including the IG Reels interest-judge DPO work.
Adds the missing `from dataclasses import dataclass` and `from typing import Optional, Tuple, TypeVar`.
___
Differential Revision: D108715300
|
@wjh70301-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D108715300. |
ErenAta16
left a comment
There was a problem hiding this comment.
This is worse than a lint fix and I think the description undersells it. The four names are used at module scope, not inside a function, so the file cannot be imported at all:
# torchtune/rlhf/loss/dpo.py on main, line 14
T = TypeVar("T", bound=dataclass)with the imports being exactly torch, torch.nn, torch.nn.functional, ChosenRejectedOutputs and deprecated. Neither TypeVar nor dataclass is among them, and that line runs at import time.
It is not a leaf module either. torchtune/rlhf/loss/__init__.py does:
from .dpo import DPOLoss, RSOLossso import torchtune.rlhf.loss raises NameError: name 'TypeVar' is not defined before anything else happens, and every consumer of DPOLoss, RSOLoss or PPOLoss goes with it. Worth saying plainly in the PR title or body, because "missing imports" reads like a typing cleanup and reviewers triage it accordingly.
The history suggests how it slipped through. 0a08c2f5 ("refactor: migrate type hints to PEP 585") removed the typing imports when it moved annotations to builtin generics, and the file still carries both spellings today, Tuple[torch.Tensor, ...] on one method and tuple[torch.Tensor, ...] on others. The later a9c32c09 ("Reference-free DPO losses", #2465) then reintroduced Optional[T], Tuple[...] and the TypeVar line without bringing the imports back. Nothing to fix there, it just explains the shape.
One thing I would change while you are in here. bound=dataclass is not a meaningful bound. dataclass is the decorator function, not a type:
T = TypeVar("T", bound=dataclass)
T.__bound__ # <function dataclass at 0x...>
isinstance(dataclass, type) # FalseIt constructs without complaint at runtime, so this will not break anything, but a type checker cannot do anything useful with a function as a bound and the annotation is effectively unbounded. If the intent is "any dataclass instance", a Protocol with __dataclass_fields__ expresses that, and if the intent was just "the concrete output type", dropping the bound entirely is more honest than one that does not check.
Since the whole point of this PR is to make the module importable, and adding from dataclasses import dataclass is what makes that bound resolvable, it seems worth settling now rather than importing a symbol solely to satisfy a line that does not do what it looks like it does.
Given the package is currently unimportable, this is worth landing quickly either way.
Summary:
torchtune/rlhf/loss/dpo.pyreferenceddataclassandTypeVar(lineT = TypeVar("T", bound=dataclass)) without importing them, raisingNameError: name 'TypeVar' is not definedat import time. This blocked every Mitra DPO unit that importstorchtune.rlhf.loss.DPOLoss— including the IG Reels interest-judge DPO work.Adds the missing
from dataclasses import dataclassandfrom typing import Optional, Tuple, TypeVar.Differential Revision: D108715300