[PyTorch] Fix deferred initialization in fusible ops - #3327
Conversation
Ops replace their params when materializing them on the first fuser forward pass, which invalidated the fuser's cached params and ops.Linear's parameter aliases. Backward then failed with a meta-vs-cuda device mismatch. - Re-cache the basic ops' params after first-forward materialization - Run first-forward initialization on the top-level ops, before the cached params are used to pick which ops require backward - Sync ops.Linear weight/bias aliases with its basic ops - Add deferred initialization tests for the affected ops Fixes NVIDIA#3322 Signed-off-by: Alp Dener <adener@nvidia.com>
Greptile SummaryThe PR fixes deferred initialization for PyTorch fusible operations by initializing top-level operations before backward selection, refreshing cached basic-operation parameters after materialization, and synchronizing
Confidence Score: 5/5The PR appears safe to merge, with no concrete blocking or non-blocking defects identified in the changed paths. First-forward initialization now reaches nested basic operations, synchronizes wrapper aliases, and refreshes the parameter cache before backward eligibility and fusion planning consume it. Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant Fuser as OperationFuser
participant Basic as Basic operations
participant Top as Top-level operations
participant Autograd
User->>Fuser: First forward
Fuser->>Basic: reset_recipe_state(recipe)
Fuser->>Top: pre_first_fuser_forward()
Top->>Basic: materialize meta parameters
Top->>Top: synchronize parameter aliases
Fuser->>Basic: cache materialized parameters
Fuser->>Fuser: select and build forward/backward fusions
Fuser->>Autograd: execute with refreshed parameters
Autograd-->>User: outputs and parameter gradients
Reviews (1): Last reviewed commit: "[PyTorch] Fix deferred initialization in..." | Re-trigger Greptile |
|
/te-ci pytorch |
vthumbe1503
left a comment
vthumbe1503
left a comment
There was a problem hiding this comment.
LGTM, have some question and a minor suggestion for adding grouped_linear test case as well for quantized_model_int case
| # Construct operation on meta device | ||
| recipe = make_recipe(quantization) | ||
| with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): | ||
| op = te_ops.Linear(size, size, bias=True, device="meta", dtype=dtype) |
There was a problem hiding this comment.
Should we extend this test for GroupedLinear as well?
| op.reset_recipe_state(recipe=recipe) | ||
| for op in self._ops: | ||
| op.pre_first_fuser_forward() | ||
|
|
||
| # Cache params now that ops have initialized them |
There was a problem hiding this comment.
If somebody calls register_parameter on the ops in a later iteration, this means it would not get cached in fuser's parameter list
I know it is a weird use-case, and maybe nobody does this. But was this always intended, @timmoon10 ?
| if op_type.startswith("grouped_linear"): | ||
| in_shape = (size * num_groups, size) | ||
| extra_inputs.append(torch.tensor([size] * num_groups, dtype=torch.int, device=device)) | ||
| x = torch.randn(in_shape, dtype=dtype, device=device, requires_grad=True) |
There was a problem hiding this comment.
We should test both requires_grad = True and False
| # Check that params have been materialized | ||
| assert op.weight is op.basic_ops[op._linear_idx].weight | ||
| assert op.bias is op.basic_ops[op._bias_idx].bias | ||
| if quantized_weight: | ||
| assert isinstance(op.weight, QuantizedTensor) | ||
| for name, param in op.named_parameters(): | ||
| assert param.device.type == device, f"{name} was not materialized on {device}" | ||
| assert param.grad is not None, f"{name} did not get a grad" | ||
| assert param.grad.device.type == device, f"{name} got a grad on {param.grad.device}" | ||
| assert x.grad is not None and x.grad.device.type == device |
There was a problem hiding this comment.
In addition to that there should be some testing that the resulting output/gradients actually match the op that did not go through this materialization step.
3 participants
Description
Ops replace their params when materializing them on the first fuser forward pass, which invalidated the fuser's cached params and ops.Linear's parameter aliases. Backward then failed with a meta-vs-cuda device mismatch.
te.ops.Linearweight/bias aliases with its basic opsFixes #3322
Type of change
Checklist: