From 2d313779979b0fd704c18627a2ea63bad393cac8 Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Fri, 3 Jul 2026 11:24:29 +0000 Subject: [PATCH] No public description PiperOrigin-RevId: 942075396 --- parallax/offload.py | 10 +++++----- parallax/offload_test.py | 4 ++-- parallax/sharding/auto_shard.py | 2 +- parallax/sharding/base_test.py | 6 +++--- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/parallax/offload.py b/parallax/offload.py index e0b14f3..0752510 100644 --- a/parallax/offload.py +++ b/parallax/offload.py @@ -104,7 +104,7 @@ def forward(model_chunk, intermediate_result): # example element 0 is the inputs to layer 0. The last element of the array # are the final outputs (which become the inputs to the loss function). saved_intermediates = [inputs] - for chunk in model.layers: + for chunk in model.layers: # pyrefly: ignore[missing-attribute] logits = forward(chunk, saved_intermediates[-1]) saved_intermediates.append(logits) @@ -132,14 +132,14 @@ def offload_backward( A full nnx.State mapping representing a composite of the grads from all layers. """ - if len(saved_intermediates) != len(model.layers) + 1: + if len(saved_intermediates) != len(model.layers) + 1: # pyrefly: ignore[missing-attribute] raise ValueError( 'The length of `saved_intermediates` must match the number of model ' 'layers plus one.' ) layer_grads = {} - for i, chunk in reversed(list(enumerate(model.layers))): + for i, chunk in reversed(list(enumerate(model.layers))): # pyrefly: ignore[missing-attribute] # pylint: disable=cell-var-from-loop chunk_input = saved_intermediates[i] graphdef, state = nnx.split(chunk) @@ -189,13 +189,13 @@ def remat_model(model: nnx.Module) -> nnx.Module: """Takes an NNX Module and returns one with all layers rematerialized.""" # TODO(jeffcarp): Generalize this to work with non-Sequential models. new_model = nnx.clone(model) - for i, layer in enumerate(new_model.layers): + for i, layer in enumerate(new_model.layers): # pyrefly: ignore[missing-attribute] signature = inspect.signature(layer.__call__) unbound_call = layer.__class__.__call__ # Shift static_argnums by 1 because of `self` in unbound call. static_argnums = tuple(i + 1 for i in range(1, len(signature.parameters))) rematted_call = nnx.remat(unbound_call, static_argnums=static_argnums) - new_model.layers[i] = RemattedLayer(layer, rematted_call) + new_model.layers[i] = RemattedLayer(layer, rematted_call) # pyrefly: ignore[missing-attribute] return new_model diff --git a/parallax/offload_test.py b/parallax/offload_test.py index fb22f32..aabd0f8 100644 --- a/parallax/offload_test.py +++ b/parallax/offload_test.py @@ -100,8 +100,8 @@ def loss_fn(model, inputs): np.testing.assert_array_equal(act_loss, exp_loss) # Assert all model weights match after gradient update. np.testing.assert_allclose( - actual_model.layers[0].kernel[...], - reference_model.layers[0].kernel[...], + actual_model.layers[0].kernel[...], # pyrefly: ignore[missing-attribute] + reference_model.layers[0].kernel[...], # pyrefly: ignore[missing-attribute] atol=1e-3, rtol=1e-1, ) diff --git a/parallax/sharding/auto_shard.py b/parallax/sharding/auto_shard.py index e680853..67bae60 100644 --- a/parallax/sharding/auto_shard.py +++ b/parallax/sharding/auto_shard.py @@ -90,7 +90,7 @@ def get_shardings( for i in range(var.aval.ndim - 1, -1, -1): # pytype: disable=attribute-error root = graph.get_root((var, i)) if ( - (model_axis is None) + (model_axis is None) # pyrefly: ignore[unbound-name] or ((root, model_axis) in edges) or (var.aval.shape[i] < min_shard_size) or dim_sharded diff --git a/parallax/sharding/base_test.py b/parallax/sharding/base_test.py index fa00f5d..2471787 100644 --- a/parallax/sharding/base_test.py +++ b/parallax/sharding/base_test.py @@ -53,7 +53,7 @@ def forward(params, inputs): return model.apply(params, inputs) output = forward(params, x) - np.testing.assert_array_almost_equal(output, reference_output, decimal=6) + np.testing.assert_array_almost_equal(output, reference_output, decimal=6) # pyrefly: ignore[bad-argument-type] # Verify compiled shardings. compiled = forward.lower(params, x).compile() # type: ignore @@ -84,7 +84,7 @@ def forward(params, inputs): return model.apply(params, inputs) output = forward(params, x) - np.testing.assert_array_almost_equal(output, reference_output, decimal=6) + np.testing.assert_array_almost_equal(output, reference_output, decimal=6) # pyrefly: ignore[bad-argument-type] # Verify compiled shardings. compiled = forward.lower(params, x).compile() # type: ignore @@ -115,7 +115,7 @@ def forward(params, inputs): sharded_forward = base.jit(forward, strategy=base.ShardingStrategy.DDP) output = sharded_forward(params, x) - np.testing.assert_array_almost_equal(output, reference_output, decimal=6) + np.testing.assert_array_almost_equal(output, reference_output, decimal=6) # pyrefly: ignore[bad-argument-type] # Verify compiled shardings. compiled = sharded_forward.lower(params, x).compile() # type: ignore