Merge branch 'main' into pre-commit-ci-update-config

Author: Borda

pyproject.toml

@@ -152,7 +152,6 @@ select = [
     "RUF100",  # see: https://docs.astral.sh/ruff/rules/unused-noqa/
 ]
 ignore = [
-    "E731",  # Do not assign a lambda expression, use a def
     "E501",  # todo: Line too long (235 > 120 characters)
     # TODO(crcrpar): Resolves the following ignores as these are added while enabling ruff check in pre-commit
     "F821",  # https://docs.astral.sh/ruff/rules/undefined-name/

thunder/benchmarks/__init__.py

@@ -44,9 +44,15 @@
 def list_benchmarks(use_classname: bool = True) -> None:
     print("Available benchmarks:")
 
-    name_fn = lambda x: x[0]
-    if not use_classname:
-        name_fn = lambda x: x[1].name
+    if use_classname:
+
+        def name_fn(x):
+            return x[0]
+
+    else:
+
+        def name_fn(x):
+            return x[1].name
 
     for x in sorted(benchmarks, key=name_fn):
         name = name_fn(x)
@@ -486,7 +492,10 @@ def _run_benchmark(
     # Determines the "wait for computation function," to be run after calls to make_batch() and the benchmark
     #   function to ensure that computation has finished
     devices: list[str] = benchmark.devices
-    wait_for_computation_fn = lambda: None
+
+    def wait_for_computation_fn():
+        return None
+
     for device in devices:
         device: thunder.core.devices.Device = thunder.core.devices.device_from_string(device)
         if device.devicetype is thunder.core.devices.DeviceType.CUDA:

thunder/benchmarks/benchmark_litgpt.py

@@ -621,7 +621,9 @@ def setup_distributed(self, model):
                 size_auto_wrap_policy = functools.partial(
                     size_based_auto_wrap_policy, min_num_params=self.fsdp_bucket_params
                 )
-                zero_bucket_wrap_policy = lambda module, recurse, nonwrapped_numel: nonwrapped_numel >= 0
+
+                def zero_bucket_wrap_policy(module, recurse, nonwrapped_numel):
+                    return nonwrapped_numel >= 0
 
                 self.bucketing_mode = self.bucketing_mode or "block"
                 custom_wrap_policy = {
@@ -657,7 +659,9 @@ def setup_activation_checkpointing(self):
             )
             return
 
-        check_fn = lambda submodule: isinstance(submodule, Block)
+        def check_fn(submodule):
+            return isinstance(submodule, Block)
+
         apply_activation_checkpointing(self.model, checkpoint_wrapper_fn=checkpoint_wrapper, check_fn=check_fn)
 
     # TODO(crcrpar): Think of apply `torch.compile` or `thunder.jit` per block/module
@@ -737,10 +741,13 @@ def calculate_model_flops(self):
             meta_model = self.init_model()
 
             x = torch.randint(0, 1, (self.micro_batch_size, meta_model.config.block_size), device=meta)
-            model_fwd = lambda: meta_model(x)
-            model_loss = lambda y: torch.nn.functional.cross_entropy(
-                y.reshape(-1, y.size(-1)), x.reshape(-1), ignore_index=-1
-            )
+
+            def model_fwd():
+                return meta_model(x)
+
+            def model_loss(y):
+                return torch.nn.functional.cross_entropy(y.reshape(-1, y.size(-1)), x.reshape(-1), ignore_index=-1)
+
             self.perf_metrics["model_flops"] = measure_flops(meta_model, model_fwd, model_loss)
         finally:
             self.device = device

thunder/benchmarks/targets.py

@@ -930,7 +930,10 @@ def test_hf_transformers(
     fn = executor(b.fn())
 
     if compute_type == ComputeType.TRAINING_BACKWARD:
-        return_fn = lambda *args, **kwargs: fn(*args, **kwargs).loss
+
+        def return_fn(*args, **kwargs):
+            return fn(*args, **kwargs).loss
+
     else:
         kwargs["labels"] = None
         return_fn = fn

thunder/core/proxies.py

@@ -1275,7 +1275,9 @@ def _infer_tensor_properties(
     else:
         # deferred computation of numel
         # TODO: similar to how `shape` is handled, this should be CSE or lifted for efficiency
-        _numel = lambda *args: reduce(operator.mul, _shape, 1)
+
+        def _numel(*args):
+            return reduce(operator.mul, _shape, 1)
 
     # TODO Alias rank to ndim?
     _ndim = len(_shape)
@@ -2017,11 +2019,14 @@ def tensorproxy(t: torch.Tensor, /, *, name: None | str, history: None | tuple =
             shape_pr = ProvenanceRecord(
                 PseudoInst.LOAD_ATTR, inputs=[copy.copy(history), wrap_const("shape").provenance]
             )
-            dim_pr = lambda idx: ProvenanceRecord(
-                PseudoInst.BINARY_SUBSCR, inputs=[shape_pr, wrap_const(idx).provenance]
-            )
+
+            def dim_pr(idx):
+                return ProvenanceRecord(PseudoInst.BINARY_SUBSCR, inputs=[shape_pr, wrap_const(idx).provenance])
+
         else:
-            dim_pr = lambda idx: None
+
+            def dim_pr(idx):
+                return None
 
         shape = tuple(
             IntegerProxy(

thunder/core/pytree.py

@@ -103,8 +103,12 @@ def register_pytree_node_dataclass(cls):
     def unpack(cls) -> dict:
         return {field.name: getattr(cls, field.name) for field in dataclasses.fields(cls)}
 
-    _flatten = lambda obj: tree_flatten(unpack(obj), namespace=OPTREE_NAMESPACE)
-    _unflatten = lambda spec, children: cls(**spec.unflatten(children))
+    def _flatten(obj):
+        return tree_flatten(unpack(obj), namespace=OPTREE_NAMESPACE)
+
+    def _unflatten(spec, children):
+        return cls(**spec.unflatten(children))
+
     optree.register_pytree_node(cls, _flatten, _unflatten, namespace=OPTREE_NAMESPACE)
     return cls
 

thunder/core/transforms.py

@@ -3110,7 +3110,10 @@ def _split_saved_for_backward_into_tensors_and_other(
     Returns:
         tuple[Sequence[Variable], Sequence[Variable]]: Tuple of tensors and other.
     """
-    is_tensor = lambda x: isinstance(x, TensorProxy)
+
+    def is_tensor(x):
+        return isinstance(x, TensorProxy)
+
     other, tensors = utils.partition(is_tensor, saved_for_backward)
     return tuple(tensors), tuple(other)
 

thunder/executors/fa3ex.py

@@ -76,7 +76,10 @@ def fa3_bwd_impl(
     dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
 
     # fa3 bwd requires last dim to be contiguous: https://github.com/Dao-AILab/flash-attention/issues/1109#issuecomment-2270043573
-    maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
+
+    def maybe_contiguous(x):
+        return x.contiguous() if x.stride(-1) != 1 else x
+
     dq, dk, dv = (maybe_contiguous(a) for a in (q, k, v))
 
     if softmax_scale is None:

thunder/executors/triton_crossentropy_impl.py

@@ -453,7 +453,10 @@ def forward(
         # result = torch.empty_like(indices, dtype=dtype, device=device)
         neg_logprobs = torch.empty_like(logits, dtype=buffer_dtype, device=device)
         weights_buffer = torch.empty_like(result, dtype=buffer_dtype)
-        grid = lambda opt: (logits.numel() // n_cols,)
+
+        def grid(opt):
+            return (logits.numel() // n_cols,)
+
         log_size_logits = int(math.log(math.prod(logits.shape) / n_cols))
         _forward[grid](
             logits,
@@ -512,10 +515,13 @@ def backward(ctx, dneg_logprobs):
         # run the kernel
         # neg_logprobs will be modified in place to become our gradient:
         n_cols = neg_logprobs.shape[-1]
-        grid = lambda opt: (
-            neg_logprobs.numel() // n_cols,
-            triton.cdiv(n_cols, opt["BLOCK"]),
-        )
+
+        def grid(opt):
+            return (
+                neg_logprobs.numel() // n_cols,
+                triton.cdiv(n_cols, opt["BLOCK"]),
+            )
+
         log_size_logits = int(math.log(math.prod(neg_logprobs.shape) / n_cols))
         _backward[grid](
             neg_logprobs,

thunder/tests/distributed/test_tensor_parallel.py

@@ -190,7 +190,10 @@ def forward(self, x):
             for p_name, p_ref in layer.named_parameters(recurse=False):
                 param_fqn = f"{l_name}.{p_name}"
                 ref_grad = p_ref.grad
-                msg = lambda err_msg: f"[{prefix} {param_fqn}] {err_msg}"
+
+                def msg(err_msg):
+                    return f"[{prefix} {param_fqn}] {err_msg}"
+
                 if is_tensor_parallel and (ref_grad.ndim > 1 or dim == 0):
                     ref_grad = ref_grad.chunk(self.world_size, dim)[self.rank]
                 grad = tp_model.get_parameter(param_fqn).grad