From 6981d338e0aa49dac502041785f8789606ef57df Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 09:00:51 +0000 Subject: [PATCH 01/10] add more nnx and sl tutorials --- .../v1/device/08_initialize_nnx_on_device.py | 485 +++++++++ .../v1/device/09_split_learning_vertical.py | 972 ++++++++++++++++++ 2 files changed, 1457 insertions(+) create mode 100644 tutorials/v1/device/08_initialize_nnx_on_device.py create mode 100644 tutorials/v1/device/09_split_learning_vertical.py diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py new file mode 100644 index 00000000..c6088a1d --- /dev/null +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -0,0 +1,485 @@ +# Copyright 2025 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tutorial 08: Initialize NNX Models on Devices and Pass State as Dicts + +Demonstrates the proper pattern for using Flax NNX with MPLang v1: +1. Initialize model on device and return state as pure Python dict +2. Pass the state dict across @mp.function boundaries +3. Reconstruct model on device using the state dict for inference +4. Train models with optimizer state management + +Key insight: NNX State can be converted to/from pure Python dicts using: +- state.to_pure_dict() → returns dict with format {layer: {param: [array, None]}} +- state.replace_by_pure_dict(dict) → reconstructs state from this format + +Optimizer state can also be managed as pure Python dicts, enabling: +- Stateful training across @mp.function boundaries +- Checkpoint-style model updates +- No need to store graphdef (can be reconstructed from model class) + +Usage: + uv run python tutorials/v1/device/08_initialize_nnx_on_device.py +""" + +import jax +import jax.numpy as jnp + +import mplang.v1 as mp +import optax +from flax import nnx + + +class SimpleMLP(nnx.Module): + """Simple Multi-Layer Perceptron for demonstration.""" + + def __init__( + self, input_dim: int, hidden_dim: int, output_dim: int, *, rngs: nnx.Rngs + ): + self.linear1 = nnx.Linear(input_dim, hidden_dim, rngs=rngs) + self.linear2 = nnx.Linear(hidden_dim, output_dim, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + x = self.linear1(x) + x = nnx.relu(x) + x = self.linear2(x) + return x + + +# Cluster configuration +cluster_spec = mp.ClusterSpec.from_dict( + { + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, + }, + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + }, + } +) + + +@mp.function +def initialize_model_and_return_state_dict( + input_dim: int, hidden_dim: int, output_dim: int, seed: int +): + """Initialize model on device P0 and return state as a pure dict. + + Returns: + Pure Python dict containing model parameters (can cross @mp.function boundaries) + """ + + def _init(): + # Create model on device + model = SimpleMLP( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + + # Split and convert state to pure dict + graphdef, state = nnx.split(model) + state_dict = state.to_pure_dict() + + print(f"[Device P0] Initialized model with {len(state_dict)} parameter groups") + print(f"[Device P0] State dict keys: {list(state_dict.keys())}") + + # Return the state dict - it's in format {layer: {param: [array, None]}} + return state_dict + + return mp.device("P0", fe_type="nnx")(_init)() + + +@mp.function +def run_inference_with_state_dict( + test_input: jax.Array, + state_dict: dict, + input_dim: int, + hidden_dim: int, + output_dim: int, + seed: int, +): + """Run inference using a state dict passed from outside. + + Args: + test_input: Input data [batch_size, input_dim] + state_dict: Pure Python dict containing model parameters + input_dim: Model input dimension + hidden_dim: Model hidden dimension + output_dim: Model output dimension + seed: Random seed (for GraphDef reconstruction) + + Returns: + Model output logits + """ + + def _infer(x, params_dict): + """Inner function that takes explicit inputs.""" + # Create an abstract model to get the GraphDef (memory efficient!) + # This only creates the structure without allocating actual arrays + abs_model = nnx.eval_shape( + lambda: SimpleMLP( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + ) + graphdef, abs_state = nnx.split(abs_model) + + print(f"[Device P0] Received state dict with keys: {list(params_dict.keys())}") + + # The params_dict is already in the format from to_pure_dict(): {key: [array, None]} + # So we can use it directly with replace_by_pure_dict + abs_state.replace_by_pure_dict(params_dict) + + print(f"[Device P0] Reconstructed state from dict") + + # Merge to get working model with the passed parameters + model = nnx.merge(graphdef, abs_state) + + print(f"[Device P0] Reconstructed model from state dict") + print(f"[Device P0] Running inference...") + + # Run inference + output = model(x) + + print(f"[Device P0] Inference complete!") + print(f" Input shape: {x.shape}") + print(f" Output shape: {output.shape}") + + return output + + return mp.device("P0", fe_type="nnx")(_infer)(test_input, state_dict) + + +@mp.function +def initialize_model_with_optimizer( + input_dim: int, hidden_dim: int, output_dim: int, seed: int, learning_rate: float +): + """Initialize model + optimizer on device P0 and return both states as dict. + + Returns: + Dict with model_state_dict and opt_state (both as pure Python dicts) + """ + + def _init(): + # Create model on device + model = SimpleMLP( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + + # Split and convert model state to pure dict + graphdef, state = nnx.split(model) + model_state_dict = state.to_pure_dict() + + # Initialize optimizer and convert its state to pure dict + # Note: optax states are already pytrees, but we'll store them explicitly + tx = optax.sgd(learning_rate) + opt_state = tx.init(model_state_dict) + + print(f"[Device P0] Initialized model with optimizer") + print(f"[Device P0] Model state keys: {list(model_state_dict.keys())}") + print(f"[Device P0] Optimizer state type: {type(opt_state)}") + + return { + "model_state_dict": model_state_dict, + "opt_state": opt_state, + "step": 0, + } + + return mp.device("P0", fe_type="nnx")(_init)() + + +@mp.function +def train_step( + train_dict: dict, + x: jnp.ndarray, + y: jnp.ndarray, + input_dim: int, + hidden_dim: int, + output_dim: int, + seed: int, + learning_rate: float, +): + """Perform one training step: forward → loss → backward → update. + + Args: + train_dict: Dict with model_state_dict, opt_state, step + x: Input batch (batch_size, input_dim) + y: Target batch (batch_size, output_dim) + input_dim: Model input dimension + hidden_dim: Model hidden dimension + output_dim: Model output dimension + seed: Random seed (for GraphDef reconstruction) + learning_rate: Learning rate for optimizer + + Returns: + Tuple of (updated_train_dict, loss_value) + """ + + def _train(train_state, batch_x, batch_y): + # Extract state + model_state_dict = train_state["model_state_dict"] + opt_state = train_state["opt_state"] + step = train_state["step"] + + # Define loss function that works with jax.grad + def loss_fn(state_dict, x, y): + # Reconstruct model from state dict for forward pass + # Create abstract model to get GraphDef (memory efficient!) + abs_model = nnx.eval_shape( + lambda: SimpleMLP( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + ) + graphdef, abs_state = nnx.split(abs_model) + abs_state.replace_by_pure_dict(state_dict) + model = nnx.merge(graphdef, abs_state) + + # Forward pass + logits = model(x) + + # MSE loss + loss = jnp.mean((logits - y) ** 2) + return loss + + # Compute gradients + loss_value = loss_fn(model_state_dict, batch_x, batch_y) + grad_fn = jax.grad(loss_fn) + grads_dict = grad_fn(model_state_dict, batch_x, batch_y) + + # Apply optimizer update + tx = optax.sgd(learning_rate) + updates, new_opt_state = tx.update(grads_dict, opt_state, model_state_dict) + new_model_state_dict = optax.apply_updates(model_state_dict, updates) + + # Return updated state + new_train_dict = { + "model_state_dict": new_model_state_dict, + "opt_state": new_opt_state, + "step": step + 1, + } + + return new_train_dict, loss_value + + return mp.device("P0", fe_type="nnx")(_train)(train_dict, x, y) + + +def main(): + """Main demonstration: initialize model on P0, then run inference using state dict.""" + print("=" * 80) + print("NNX Model with State Dict Pattern") + print("=" * 80) + + # Setup simulator + simulator = mp.Simulator(cluster_spec) + + # Step 1: Initialize model on device P0 and get state as dict + print("\n[Step 1] Initializing model on device P0...") + state_dict_result = mp.evaluate( + simulator, + initialize_model_and_return_state_dict, + input_dim=8, + hidden_dim=16, + output_dim=2, + seed=42, + ) + + # Fetch the state dict to driver + state_dict = mp.fetch(simulator, state_dict_result) + print(f"\n[Driver] Received state dict with keys: {list(state_dict.keys())}") + + # Step 2: Run inference using the state dict + print("\n[Step 2] Running inference with state dict...") + + # Create test input + test_key = jax.random.PRNGKey(999) + test_input = jax.random.normal(test_key, (4, 8)) + + output_result = mp.evaluate( + simulator, + run_inference_with_state_dict, + test_input, + state_dict_result, + input_dim=8, + hidden_dim=16, + output_dim=2, + seed=42, # Same seed to get same GraphDef + ) + + output = mp.fetch(simulator, output_result) + + # Handle output if it's wrapped + if isinstance(output, (list, tuple)) and len(output) > 0: + output = output[0] + if not isinstance(output, jnp.ndarray): + output = jnp.array(output) + + print(f"\n[Driver] Inference complete!") + print(f" Output shape: {output.shape}") + print(f" Output sample: {output[0]}") + + # Step 3: Initialize model with optimizer + print("\n[Step 3] Initializing model with optimizer on device P0...") + train_state_dict_result = mp.evaluate( + simulator, + initialize_model_with_optimizer, + input_dim=8, + hidden_dim=16, + output_dim=2, + seed=42, + learning_rate=0.01, + ) + + # Fetch the train state dict to driver + train_state_dict = mp.fetch(simulator, train_state_dict_result) + print( + f"\n[Driver] Received train state dict with keys: {list(train_state_dict.keys())}" + ) + + # Step 4: Perform a single training step + print("\n[Step 4] Performing a single training step...") + + # Create a simple learnable dataset: y = 2*x + noise + # This gives the model something real to learn + x_train = jax.random.normal(jax.random.PRNGKey(888), (32, 8)) + # Create target with a simple linear relationship + true_weights = jnp.array( + [[2.0], [-1.0], [0.5], [1.5], [-0.5], [1.0], [-1.5], [0.0]] + ) + true_weights2 = jnp.array([[1.0], [-0.5]]) + y_train = ( + x_train @ true_weights @ true_weights2.T + + jax.random.normal(jax.random.PRNGKey(999), (32, 2)) * 0.1 + ) + + train_result = mp.evaluate( + simulator, + train_step, + train_state_dict_result, + x_train, + y_train, + input_dim=8, + hidden_dim=16, + output_dim=2, + seed=42, + learning_rate=0.01, + ) + + # Fetch only for human inspection + updated_train_dict, loss_value = mp.fetch(simulator, train_result) + print(f"\n[Driver] Training step complete!") + print(f" Loss value: {loss_value}") + print( + f" Updated model state keys: {list(updated_train_dict['model_state_dict'].keys())}" + ) + print(f" Updated optimizer state type: {type(updated_train_dict['opt_state'])}") + print(f" Step number: {updated_train_dict['step']}") + + # Step 5: Demonstrate Optimizer State Persists Across Multiple Training Steps + print("\n" + "=" * 80) + print("[Step 5] Optimizer State Persists Across MPLang Function Boundaries") + print("=" * 80) + print("Demonstrating that optimizer state can be passed back and forth") + print("\nRunning 5 more training steps with THE SAME batch...") + print("Note: Loss should decrease steadily as the model learns the pattern!\n") + + # Keep using train_result (before fetch) for subsequent computations + current_train_result = train_result + + for step_idx in range(5): + # THE KEY POINT: Pass train_result[0] (the train_dict) to next mp.evaluate! + # train_result is a tuple (train_dict, loss), we need just the train_dict + # Use the SAME training data to show actual learning + current_train_result = mp.evaluate( + simulator, + train_step, + current_train_result[ + 0 + ], # Pass just the train_dict (first element of tuple) + x_train, # Same data - so we can see the model actually learning! + y_train, + input_dim=8, + hidden_dim=16, + output_dim=2, + seed=42, + learning_rate=0.01, + ) + + # Fetch only for display (human inspection) + updated_train_dict, loss_value = mp.fetch(simulator, current_train_result) + print( + f" Step {step_idx + 2}: Loss = {loss_value}, Global step = {updated_train_dict['step']}" + ) + + # Final fetch for summary + final_train_dict, final_loss = mp.fetch(simulator, current_train_result) + + print("\n" + "=" * 80) + print("Training Progress Summary") + print("=" * 80) + print(f"Final loss: {final_loss}") + print(f"Final optimizer state type: {type(final_train_dict['opt_state'])}") + print(f"Final global step: {final_train_dict['step']}") + + print("\n" + "=" * 80) + print("Key Takeaways") + print("=" * 80) + print("1. ✅ Initialize NNX models on devices with mp.device('P0', fe_type='nnx')") + print("2. ✅ Convert state to pure dict: state.to_pure_dict()") + print(" - Format: {layer: {param: [array, None]}}") + print("3. ✅ Pass dict across @mp.function boundaries (works in MPLang v1!)") + print("4. ✅ Reconstruct state: state.replace_by_pure_dict(state_dict)") + print("5. ✅ Use nnx.eval_shape() for memory-efficient GraphDef creation") + print(" - Creates abstract model without allocating arrays!") + print("6. ✅ Merge with GraphDef: model = nnx.merge(graphdef, state)") + print("7. ✅ This enables efficient model reuse without reinitializing weights!") + print("8. ✅ Initialize model with optimizer state management") + print(" - Use optax for flexible optimizer integration") + print(" - Manage optimizer state as pure Python dicts") + print( + "9. ✅ Perform training steps with full control: forward, loss, backward, update" + ) + print(" - Stateful training across @mp.function boundaries") + print(" - Checkpoint-style updates with pure dicts") + print( + "10. ✅ CRITICAL: Optimizer state persists across MPLang function boundaries!" + ) + print(" - Pass train_dict (with opt_state) between training steps") + print(" - Optimizer accumulates state (momentum, etc.) correctly") + print(" - Step counter tracks global training progress") + print(" - Enables iterative training with stateful optimizers") + print(" - Loss decreases steadily when training on same batch!") + print("\nThis pattern solves the MPLang v1 limitation of passing complex types") + print("by using pure Python dicts that can cross @mp.function boundaries.") + print("Optimizer state management enables proper multi-step training!") + print("=" * 80) + + +if __name__ == "__main__": + main() diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py new file mode 100644 index 00000000..dffcf6e4 --- /dev/null +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -0,0 +1,972 @@ +# Copyright 2025 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tutorial 09: Split Learning with Vertical Data Partitioning + +Demonstrates split learning in MPLang v1 where: +- **Data is vertically partitioned**: Alice (P0) has n samples × m1 features + labels, + Bob (P1) has n samples × m2 features (no labels) +- **Three-model architecture**: Alice base model, Bob base model, Alice aggregate model +- **Privacy-preserving**: Raw features never leave their owner's device +- **Training flow**: Forward → backward → update using surrogate loss for base models + +Architecture: + Alice (P0): + - Alice Base Model: m1 features → h1 embeddings (m1 → 16 → 16) + - Alice Aggregate Model: [h1, h2] → logits (32 → 16 → 2) + - Has labels, computes loss + + Bob (P1): + - Bob Base Model: m2 features → h2 embeddings (m2 → 16 → 16) + - No labels, uses surrogate loss + +Key Concepts: + 1. **State Management**: Use graphdef + state dict pattern (following Tutorial 08) + 2. **Memory Efficiency**: Use nnx.eval_shape() to create abstract models (no array allocation) + 3. **Surrogate Loss**: Base models use dot product (grad · embedding) for training + 4. **Privacy**: Only embeddings and gradients are shared, not raw features + 5. **Optimizer**: Recreated on each step (SGD with fixed learning rate) + +Training Flow (One Iteration): + 1. Alice forward: x_alice → h1 + 2. Bob forward: x_bob → h2 + 3. Transfer h2 from P1 to P0 + 4. Alice aggregate: [h1, h2] → logits → loss (supervised) + 5. Alice aggregate backward: compute gradients (params, h1, h2) + 6. Alice aggregate update: apply gradients + 7. Alice base backward: compute gradients using surrogate loss (grad_h1 · h1) + 8. Alice base update: apply gradients + 9. Transfer grad_h2 from P0 to P1 + 10. Bob base backward: compute gradients using surrogate loss (grad_h2 · h2) + 11. Bob base update: apply gradients + +Usage: + uv run python tutorials/v1/device/09_split_learning_vertical.py +""" + +import os +from functools import partial + +import jax +import jax.numpy as jnp + +import mplang.v1 as mp +import optax +import pandas as pd +from flax import nnx +from mplang.v1.core.dtypes import FLOAT64, INT64 +from mplang.v1.ops import basic as basic_ops + +# We'll use a simplified state management approach similar to Tutorial 08 +# Instead of TrainState, we'll manually manage state dict + optimizer state + +# ============================================================================ +# Configuration +# ============================================================================ + +# Dataset parameters +N_SAMPLES = 10000 # Number of samples +M1 = 10 # Number of Alice's features +M2 = 8 # Number of Bob's features +H1 = H2 = 16 # Embedding dimensions (h1 = h2 = 16) +N_CLASSES = 2 # Binary classification + +# Training parameters +LEARNING_RATE = 0.01 +SEED_ALICE = 42 +SEED_BOB = 43 +SEED_AGG = 44 + +# Cluster specification +cluster_spec = mp.ClusterSpec.from_dict( + { + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, + }, + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, + }, + } +) + + +# ============================================================================ +# Model Definitions (Flax NNX) +# ============================================================================ + + +class AliceBaseModel(nnx.Module): + """Alice's base model: m1 features → h1 embeddings (m1 → 16 → 16).""" + + def __init__(self, input_dim: int, hidden_dim: int, *, rngs: nnx.Rngs): + """Initialize Alice's base model. + + Args: + input_dim: Number of input features (m1) + hidden_dim: Embedding dimension (h1) + rngs: Random number generator state + """ + self.linear1 = nnx.Linear(input_dim, hidden_dim, rngs=rngs) + self.linear2 = nnx.Linear(hidden_dim, hidden_dim, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + """Forward pass: x (n, m1) → h1 (n, h1).""" + x = self.linear1(x) + x = nnx.relu(x) + x = self.linear2(x) + return x + + +class BobBaseModel(nnx.Module): + """Bob's base model: m2 features → h2 embeddings (m2 → 16 → 16).""" + + def __init__(self, input_dim: int, hidden_dim: int, *, rngs: nnx.Rngs): + """Initialize Bob's base model. + + Args: + input_dim: Number of input features (m2) + hidden_dim: Embedding dimension (h2) + rngs: Random number generator state + """ + self.linear1 = nnx.Linear(input_dim, hidden_dim, rngs=rngs) + self.linear2 = nnx.Linear(hidden_dim, hidden_dim, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + """Forward pass: x (n, m2) → h2 (n, h2).""" + x = self.linear1(x) + x = nnx.relu(x) + x = self.linear2(x) + return x + + +class AliceAggregateModel(nnx.Module): + """Alice's aggregate model: [h1, h2] → logits (32 → 16 → 2).""" + + def __init__( + self, input_dim: int, hidden_dim: int, output_dim: int, *, rngs: nnx.Rngs + ): + """Initialize Alice's aggregate model. + + Args: + input_dim: Combined embedding dimension (h1 + h2 = 32) + hidden_dim: Hidden layer dimension (16) + output_dim: Number of classes (2) + rngs: Random number generator state + """ + self.linear1 = nnx.Linear(input_dim, hidden_dim, rngs=rngs) + self.linear2 = nnx.Linear(hidden_dim, output_dim, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + """Forward pass: x (n, 32) → logits (n, 2).""" + x = self.linear1(x) + x = nnx.relu(x) + x = self.linear2(x) + return x + + +# ============================================================================ +# Helper Functions: State Management +# ============================================================================ + + +def model_state_to_dict(state, opt_state, step): + """Convert model state to pure Python dict for cross-device transfer. + + Note: graphdef is NOT stored - it can be reconstructed from model class + + Args: + state: NNX State (model parameters) + opt_state: Optax optimizer state + step: Training step counter + + Returns: + Dict with keys: model_state_dict, opt_state, step + """ + return { + "model_state_dict": state.to_pure_dict(), + "opt_state": opt_state, + "step": step, + } + + +def reconstruct_model_from_dict( + model_dict: dict, model_class, *model_args, **model_kwargs +): + """Reconstruct model and state from dict. + + GraphDef is reconstructed from model class definition. + + Args: + model_dict: Dict with keys: model_state_dict, opt_state, step + model_class: Model class (AliceBaseModel, BobBaseModel, AliceAggregateModel) + *model_args: Arguments for model initialization + **model_kwargs: Keyword arguments for model initialization + + Returns: + Tuple of (graphdef, state, opt_state, step) + """ + # Create temporary model to get GraphDef structure + temp_model = model_class(*model_args, **model_kwargs) + graphdef, temp_state = nnx.split(temp_model) + + # Replace with actual parameters + temp_state.replace_by_pure_dict(model_dict["model_state_dict"]) + + return graphdef, temp_state, model_dict["opt_state"], model_dict["step"] + + +# ============================================================================ +# Helper Functions: Forward Pass +# ============================================================================ + + +def alice_base_forward(x, model_dict, m1, h1, seed): + """Alice's base model forward pass. + + Args: + x: Input features (n, m1) + model_dict: Alice base model state as dict + m1: Number of Alice's features + h1: Embedding dimension + seed: Random seed for model reconstruction + + Returns: + h1: Alice's embeddings (n, h1) + """ + # Reconstruct model + graphdef, state, _, _ = reconstruct_model_from_dict( + model_dict, AliceBaseModel, m1, h1, rngs=nnx.Rngs(seed) + ) + model = nnx.merge(graphdef, state) + + # Forward pass + h1 = model(x) + + return h1 + + +def bob_base_forward(x, model_dict, m2, h2, seed): + """Bob's base model forward pass. + + Args: + x: Input features (n, m2) + model_dict: Bob base model state as dict + m2: Number of Bob's features + h2: Embedding dimension + seed: Random seed for model reconstruction + + Returns: + h2: Bob's embeddings (n, h2) + """ + # Reconstruct model + graphdef, state, _, _ = reconstruct_model_from_dict( + model_dict, BobBaseModel, m2, h2, rngs=nnx.Rngs(seed) + ) + model = nnx.merge(graphdef, state) + + # Forward pass + h2 = model(x) + + return h2 + + +def alice_aggregate_forward_and_loss(h1, h2, y, model_dict, n_classes, seed): + """Alice's aggregate model forward pass with loss computation. + + Args: + h1: Alice's embeddings (n, h1) + h2: Bob's embeddings (n, h2) + y: Labels (n,) + model_dict: Alice aggregate model state as dict + n_classes: Number of output classes + seed: Random seed for model reconstruction + + Returns: + loss: Scalar loss value + logits: Model predictions (n, n_classes) + combined: Concatenated embeddings (n, h1+h2) + """ + # Reconstruct model + graphdef, state, _, _ = reconstruct_model_from_dict( + model_dict, AliceAggregateModel, H1 + H2, 16, n_classes, rngs=nnx.Rngs(seed) + ) + model = nnx.merge(graphdef, state) + + # Concatenate embeddings + combined = jnp.concatenate([h1, h2], axis=1) + + # Forward pass + logits = model(combined) + + # Compute loss (cross-entropy) + loss = jnp.mean(optax.softmax_cross_entropy_with_integer_labels(logits, y)) + + return loss, logits, combined + + +# ============================================================================ +# Helper Functions: Backward Pass +# ============================================================================ + + +def alice_aggregate_backward(h1, h2, y, model_dict, n_classes, seed): + """Compute gradients for Alice's aggregate model using actual loss. + + This is standard supervised learning - compute loss from predictions and labels. + + Args: + h1: Alice's embeddings (n, h1) + h2: Bob's embeddings (n, h2) + y: Labels (n,) + model_dict: Alice aggregate model state as dict + n_classes: Number of output classes + seed: Random seed for model reconstruction + + Returns: + grads_state: Gradients for aggregate model parameters + grad_h1: Gradient w.r.t. Alice's embeddings (n, h1) + grad_h2: Gradient w.r.t. Bob's embeddings (n, h2) + loss: Scalar loss value + """ + + # Define loss function + def loss_fn(state_dict, h1, h2, y): + # Reconstruct model from state dict + # n_classes and seed are captured from outer scope and are concrete values + temp_model = AliceAggregateModel(H1 + H2, 16, n_classes, rngs=nnx.Rngs(seed)) + graphdef, temp_state = nnx.split(temp_model) + temp_state.replace_by_pure_dict(state_dict) + model = nnx.merge(graphdef, temp_state) + + # Forward pass + combined = jnp.concatenate([h1, h2], axis=1) + logits = model(combined) + loss = jnp.mean(optax.softmax_cross_entropy_with_integer_labels(logits, y)) + return loss + + # Get model state dict + model_state_dict = model_dict["model_state_dict"] + + # Compute gradients w.r.t. state_dict and inputs (h1, h2) + grad_fn = jax.grad(loss_fn, argnums=(0, 1, 2)) + grads_state_dict, grad_h1, grad_h2 = grad_fn(model_state_dict, h1, h2, y) + + # Also compute loss for logging + loss = loss_fn(model_state_dict, h1, h2, y) + + return grads_state_dict, grad_h1, grad_h2, loss + + +def alice_base_backward(x_alice, grad_h1, model_dict, m1, h1, seed): + """Compute gradients for Alice's base model using surrogate loss. + + Split Learning Insight: + - Alice's base model doesn't have direct access to labels + - Use surrogate loss: L_surrogate = grad_h1 · h1 (dot product) + - This mimics backpropagation from the aggregate model + - grad_h1 comes from ∂L_aggregate/∂h1 + + Args: + x_alice: Input features (n, m1) + grad_h1: Gradient from aggregate model (n, h1) + model_dict: Alice base model state as dict + m1: Number of Alice's features + h1: Embedding dimension + seed: Random seed for model reconstruction + + Returns: + grads_state_dict: Gradients for base model parameters (as dict) + """ + + # Surrogate loss function + def surrogate_loss_fn(state_dict, x, grad_from_next_layer): + # Reconstruct model from state dict + # m1, h1, seed are captured from outer scope and are concrete values + temp_model = AliceBaseModel(m1, h1, rngs=nnx.Rngs(seed)) + graphdef, temp_state = nnx.split(temp_model) + temp_state.replace_by_pure_dict(state_dict) + model = nnx.merge(graphdef, temp_state) + + h = model(x) + # Surrogate loss: dot product with gradients from next layer + loss = jnp.sum(h * grad_from_next_layer) + return loss + + # Get model state dict + model_state_dict = model_dict["model_state_dict"] + + # Compute gradients w.r.t. state_dict using surrogate loss + grad_fn = jax.grad(surrogate_loss_fn) + grads_state_dict = grad_fn(model_state_dict, x_alice, grad_h1) + + return grads_state_dict + + +def bob_base_backward(x_bob, grad_h2, model_dict, m2, h2, seed): + """Compute gradients for Bob's base model using surrogate loss. + + Split Learning Insight: + - Bob's base model doesn't have direct access to labels + - Use surrogate loss: L_surrogate = grad_h2 · h2 (dot product) + - This mimics backpropagation from the aggregate model + - grad_h2 comes from ∂L_aggregate/∂h2 + + Args: + x_bob: Input features (n, m2) + grad_h2: Gradient from aggregate model (n, h2) + model_dict: Bob base model state as dict + m2: Number of Bob's features + h2: Embedding dimension + seed: Random seed for model reconstruction + + Returns: + grads_state_dict: Gradients for base model parameters (as dict) + """ + + # Surrogate loss function + def surrogate_loss_fn(state_dict, x, grad_from_next_layer): + # Reconstruct model from state dict + # m2, h2, seed are captured from outer scope and are concrete values + temp_model = BobBaseModel(m2, h2, rngs=nnx.Rngs(seed)) + graphdef, temp_state = nnx.split(temp_model) + temp_state.replace_by_pure_dict(state_dict) + model = nnx.merge(graphdef, temp_state) + + h = model(x) + # Surrogate loss: dot product with gradients from next layer + loss = jnp.sum(h * grad_from_next_layer) + return loss + + # Get model state dict + model_state_dict = model_dict["model_state_dict"] + + # Compute gradients w.r.t. state_dict using surrogate loss + grad_fn = jax.grad(surrogate_loss_fn) + grads_state_dict = grad_fn(model_state_dict, x_bob, grad_h2) + + return grads_state_dict + + +# ============================================================================ +# Helper Functions: Update +# ============================================================================ + + +def update_model_state(model_dict, grads_dict, lr): + """Update model state with gradients using optax SGD. + + Similar to Tutorial 08 pattern - use optax for clean optimizer integration. + + Args: + model_dict: Current model state dict (keys: model_state_dict, opt_state, step) + grads_dict: Gradients as pure dict (same format as model_state_dict) + lr: Learning rate + + Returns: + new_model_dict: Updated model state dict + """ + model_state_dict = model_dict["model_state_dict"] + opt_state = model_dict["opt_state"] + step = model_dict["step"] + + # Create optimizer (SGD) + tx = optax.sgd(lr) + + # Apply optimizer update + updates, new_opt_state = tx.update(grads_dict, opt_state, model_state_dict) + new_model_state_dict = optax.apply_updates(model_state_dict, updates) + + # Return updated dict + return { + "model_state_dict": new_model_state_dict, + "opt_state": new_opt_state, + "step": step + 1, + } + + +# ============================================================================ +# Helper Functions: Reusable Partial Functions +# ============================================================================ + +# Create partial functions that bind module-level constants (not traced parameters) +# These can be reused in training steps without worrying about JAX tracer issues + + +def get_alice_base_backward_fn(): + """Returns partial function for Alice base backward with module constants.""" + return partial(alice_base_backward, m1=M1, h1=H1, seed=SEED_ALICE) + + +def get_bob_base_backward_fn(): + """Returns partial function for Bob base backward with module constants.""" + return partial(bob_base_backward, m2=M2, h2=H2, seed=SEED_BOB) + + +def get_alice_agg_backward_fn(): + """Returns partial function for Alice aggregate backward with module constants.""" + return partial(alice_aggregate_backward, n_classes=N_CLASSES, seed=SEED_AGG) + + +# ============================================================================ +# Data Preparation +# ============================================================================ + + +def prepare_vertical_split_data(): + """Generate synthetic vertically partitioned data for split learning. + + Creates: + - Alice (P0): 10,000 samples × 10 features + labels + - Bob (P1): 10,000 samples × 8 features (no labels) + + Saves to: + - tmp/alice_data.csv + - tmp/bob_data.csv + """ + print("\n" + "=" * 80) + print("Generating Synthetic Vertically Partitioned Data") + print("=" * 80) + + # Create tmp directory + os.makedirs("tmp", exist_ok=True) + base_dir = os.path.abspath("tmp") + + # Generate synthetic data + key = jax.random.PRNGKey(0) + key_alice, key_bob, key_labels = jax.random.split(key, 3) + + # Alice: m1 features + labels + alice_features = jax.random.normal(key_alice, (N_SAMPLES, M1)) + alice_labels = jax.random.randint(key_labels, (N_SAMPLES,), 0, N_CLASSES) + + # Bob: m2 features (no labels) + bob_features = jax.random.normal(key_bob, (N_SAMPLES, M2)) + + # Save as CSV + alice_cols = {f"alice_f{i}": alice_features[:, i] for i in range(M1)} + alice_cols["label"] = alice_labels + df_alice = pd.DataFrame(alice_cols) + df_alice.to_csv(f"{base_dir}/alice_data.csv", index=False) + + bob_cols = {f"bob_f{i}": bob_features[:, i] for i in range(M2)} + df_bob = pd.DataFrame(bob_cols) + df_bob.to_csv(f"{base_dir}/bob_data.csv", index=False) + + print(f"✅ Alice data: {N_SAMPLES} samples × {M1} features + labels") + print(f" Saved to: {base_dir}/alice_data.csv") + print(f"✅ Bob data: {N_SAMPLES} samples × {M2} features") + print(f" Saved to: {base_dir}/bob_data.csv") + print("=" * 80) + + +@mp.function +def load_vertical_split_data( + alice_csv: str, bob_csv: str, m1: int, m2: int, n_rows: int +): + """Load vertically split data onto P0 and P1. + + Args: + alice_csv: Path to Alice's CSV file + bob_csv: Path to Bob's CSV file + m1: Number of Alice's features + m2: Number of Bob's features + n_rows: Number of rows to load + + Returns: + alice_features: Tensor on P0 (n, m1) + alice_labels: Tensor on P0 (n,) + bob_features: Tensor on P1 (n, m2) + """ + # Define schemas (all columns must have same dtype for table_to_tensor) + # Cast label to FLOAT64 here, convert back to int after splitting + schema_alice = mp.TableType.from_dict( + {**{f"alice_f{i}": FLOAT64 for i in range(m1)}, "label": FLOAT64} + ) + schema_bob = mp.TableType.from_dict({f"bob_f{i}": FLOAT64 for i in range(m2)}) + + # Read CSVs as tables on respective devices + tbl_alice = mp.device("P0")(basic_ops.read)(path=alice_csv, ty=schema_alice) + tbl_bob = mp.device("P1")(basic_ops.read)(path=bob_csv, ty=schema_bob) + + # Convert to tensors + alice_tensor = mp.device("P0")(basic_ops.table_to_tensor)( + tbl_alice, number_rows=n_rows + ) + bob_tensor = mp.device("P1")(basic_ops.table_to_tensor)(tbl_bob, number_rows=n_rows) + + # Split Alice tensor into features and labels + def split_features_labels(data): + features = data[:, :-1] # All columns except last + labels = data[:, -1].astype(jnp.int32) # Last column as int + return features, labels + + alice_features, alice_labels = mp.device("P0")(split_features_labels)(alice_tensor) + bob_features = bob_tensor + + return alice_features, alice_labels, bob_features + + +# ============================================================================ +# Model Initialization +# ============================================================================ + + +@mp.function +def initialize_alice_base_model(m1: int, h1: int, seed: int, learning_rate: float): + """Initialize Alice's base model on P0 and return state as dict.""" + + def _init(): + # Create model + model = AliceBaseModel(input_dim=m1, hidden_dim=h1, rngs=nnx.Rngs(seed)) + + # Split into graphdef + state + graphdef, state = nnx.split(model) + + # Convert state to pure dict + model_state_dict = state.to_pure_dict() + + # Initialize optimizer + tx = optax.sgd(learning_rate) + opt_state = tx.init(model_state_dict) + + return {"model_state_dict": model_state_dict, "opt_state": opt_state, "step": 0} + + return mp.device("P0", fe_type="nnx")(_init)() + + +@mp.function +def initialize_bob_base_model(m2: int, h2: int, seed: int, learning_rate: float): + """Initialize Bob's base model on P1 and return state as dict.""" + + def _init(): + # Create model + model = BobBaseModel(input_dim=m2, hidden_dim=h2, rngs=nnx.Rngs(seed)) + + # Split into graphdef + state + graphdef, state = nnx.split(model) + + # Convert state to pure dict + model_state_dict = state.to_pure_dict() + + # Initialize optimizer + tx = optax.sgd(learning_rate) + opt_state = tx.init(model_state_dict) + + return {"model_state_dict": model_state_dict, "opt_state": opt_state, "step": 0} + + return mp.device("P1", fe_type="nnx")(_init)() + + +@mp.function +def initialize_alice_agg_model( + input_dim: int, hidden_dim: int, output_dim: int, seed: int, learning_rate: float +): + """Initialize Alice's aggregate model on P0 and return state as dict.""" + + def _init(): + # Create model + model = AliceAggregateModel( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + + # Split into graphdef + state + graphdef, state = nnx.split(model) + + # Convert state to pure dict + model_state_dict = state.to_pure_dict() + + # Initialize optimizer + tx = optax.sgd(learning_rate) + opt_state = tx.init(model_state_dict) + + return {"model_state_dict": model_state_dict, "opt_state": opt_state, "step": 0} + + return mp.device("P0", fe_type="nnx")(_init)() + + +# ============================================================================ +# Training Step +# ============================================================================ + + +@mp.function +def split_learning_train_step( + alice_features, # On P0: (10000, m1) + alice_labels, # On P0: (10000,) + bob_features, # On P1: (10000, m2) + alice_base_model_dict, # Alice base model state as dict + bob_base_model_dict, # Bob base model state as dict + alice_agg_model_dict, # Alice aggregate model state as dict + learning_rate: float, + m1: int, + m2: int, + h1: int, + h2: int, + n_classes: int, + seed_alice: int, + seed_bob: int, + seed_agg: int, +): + """One complete split learning training step (single batch = full dataset). + + Training Flow (Split Learning): + 1. Alice computes base model forward → h1 + 2. Bob computes base model forward → h2 + 3. Transfer h2 from P1 to P0 + 4. Alice runs aggregate model → logits → loss (actual supervised loss) + 5. Alice computes gradients for aggregate model (∂L/∂params_agg, ∂L/∂h1, ∂L/∂h2) + 6. Alice updates aggregate model with actual gradients + 7. Alice computes gradients for her base model using surrogate loss (grad_h1 · h1) + 8. Alice updates her base model with surrogate gradients + 9. Alice sends ∂L/∂h2 back to Bob (P0 → P1) + 10. Bob computes gradients for his base model using surrogate loss (grad_h2 · h2) + 11. Bob updates his base model with surrogate gradients + + Note: Uses helper functions (get_*_backward_fn) that bind module-level constants + via functools.partial to avoid traced parameter issues in JAX gradient computation. + + Returns: + new_alice_base_model_dict: Updated Alice base model state (as dict) + new_bob_base_model_dict: Updated Bob base model state (as dict) + new_alice_agg_model_dict: Updated Alice aggregate model state (as dict) + loss: Training loss (scalar) + """ + + # === Forward Pass === + + # 1. Alice base model forward on P0 + h1 = mp.device("P0", fe_type="nnx")(alice_base_forward)( + alice_features, alice_base_model_dict, m1, h1, seed_alice + ) + + # 2. Bob base model forward on P1 + h2_on_p1 = mp.device("P1", fe_type="nnx")(bob_base_forward)( + bob_features, bob_base_model_dict, m2, h2, seed_bob + ) + + # 3. Transfer Bob's embeddings to Alice (P1 → P0) + h2_on_p0 = mp.put("P0", h2_on_p1) + + # === Backward Pass (Alice Aggregate Model) === + + # 4. Alice computes gradients for aggregate model using actual loss + agg_grads_dict, grad_h1, grad_h2, loss = mp.device("P0", fe_type="nnx")( + get_alice_agg_backward_fn() + )(h1, h2_on_p0, alice_labels, alice_agg_model_dict) + + # 5. Alice updates her aggregate model with actual gradients + new_alice_agg_model_dict = mp.device("P0", fe_type="nnx")(update_model_state)( + alice_agg_model_dict, agg_grads_dict, learning_rate + ) + + # === Backward Pass (Alice Base Model) === + + # 6. Alice computes gradients for her base model using surrogate loss + alice_base_grads_dict = mp.device("P0", fe_type="nnx")( + get_alice_base_backward_fn() + )(alice_features, grad_h1, alice_base_model_dict) + + # 7. Alice updates her base model with surrogate gradients + new_alice_base_model_dict = mp.device("P0", fe_type="nnx")(update_model_state)( + alice_base_model_dict, alice_base_grads_dict, learning_rate + ) + + # === Backward Pass (Bob Base Model) === + + # 8. Transfer grad_h2 to Bob (P0 → P1) + grad_h2_on_p1 = mp.put("P1", grad_h2) + + # 9. Bob computes gradients for his base model using surrogate loss + bob_base_grads_dict = mp.device("P1", fe_type="nnx")(get_bob_base_backward_fn())( + bob_features, grad_h2_on_p1, bob_base_model_dict + ) + + # 10. Bob updates his base model with surrogate gradients + new_bob_base_model_dict = mp.device("P1", fe_type="nnx")(update_model_state)( + bob_base_model_dict, bob_base_grads_dict, learning_rate + ) + + return ( + new_alice_base_model_dict, + new_bob_base_model_dict, + new_alice_agg_model_dict, + loss, + ) + + +# ============================================================================ +# Main Function +# ============================================================================ + + +def main(): + """Main demonstration of split learning with vertical data partitioning.""" + print("\n" + "=" * 80) + print("Split Learning Tutorial: Vertical Data Partitioning") + print("=" * 80) + print(f"Dataset: {N_SAMPLES} samples") + print(f"Alice (P0): {M1} features + labels") + print(f"Bob (P1): {M2} features (no labels)") + print(f"Embedding dimensions: h1 = h2 = {H1}") + print(f"Number of classes: {N_CLASSES}") + print(f"Learning rate: {LEARNING_RATE}") + print("=" * 80) + + # Step 1: Prepare data + prepare_vertical_split_data() + + # Step 2: Setup simulator + print("\n[Step 2] Setting up MPLang simulator...") + simulator = mp.Simulator(cluster_spec) + + # Step 3: Load data onto devices + print("\n[Step 3] Loading data onto devices P0 and P1...") + base_dir = os.path.abspath("tmp") + alice_features, alice_labels, bob_features = mp.evaluate( + simulator, + load_vertical_split_data, + f"{base_dir}/alice_data.csv", + f"{base_dir}/bob_data.csv", + M1, + M2, + N_SAMPLES, + ) + print("✅ Data loaded successfully") + + # Step 4: Initialize models + print("\n[Step 4] Initializing models on devices...") + + # Alice base model (P0) + alice_base_model_dict = mp.evaluate( + simulator, + initialize_alice_base_model, + M1, + H1, + SEED_ALICE, + LEARNING_RATE, + ) + print(f"✅ Alice base model initialized on P0 (m1={M1} → h1={H1})") + + # Bob base model (P1) + bob_base_model_dict = mp.evaluate( + simulator, + initialize_bob_base_model, + M2, + H2, + SEED_BOB, + LEARNING_RATE, + ) + print(f"✅ Bob base model initialized on P1 (m2={M2} → h2={H2})") + + # Alice aggregate model (P0) + alice_agg_model_dict = mp.evaluate( + simulator, + initialize_alice_agg_model, + H1 + H2, # Concatenated embeddings + 16, # Hidden dimension + N_CLASSES, + SEED_AGG, + LEARNING_RATE, + ) + print(f"✅ Alice aggregate model initialized on P0 ({H1+H2} → 16 → {N_CLASSES})") + + # Step 5: Run training iterations + print("\n[Step 5] Running split learning training iterations...") + print("Training flow:") + print(" 1. Alice forward: x_alice → h1") + print(" 2. Bob forward: x_bob → h2") + print(" 3. Transfer h2: P1 → P0") + print(" 4. Alice aggregate: [h1, h2] → loss") + print(" 5. Alice aggregate backward & update") + print(" 6. Alice base backward & update (surrogate loss)") + print(" 7. Transfer grad_h2: P0 → P1") + print(" 8. Bob base backward & update (surrogate loss)") + print("\nRunning 5 training iterations on the same batch...") + + current_alice_base = alice_base_model_dict + current_bob_base = bob_base_model_dict + current_alice_agg = alice_agg_model_dict + + for iter_idx in range(5): + result = mp.evaluate( + simulator, + split_learning_train_step, + alice_features, + alice_labels, + bob_features, + current_alice_base, + current_bob_base, + current_alice_agg, + LEARNING_RATE, + M1, + M2, + H1, + H2, + N_CLASSES, + SEED_ALICE, + SEED_BOB, + SEED_AGG, + ) + + # Fetch only for inspection + ( + new_alice_base_fetched, + new_bob_base_fetched, + new_alice_agg_fetched, + loss_fetched, + ) = mp.fetch(simulator, result) + + print(f" Iteration {iter_idx + 1}: Loss = {loss_fetched}") + + # Update current state using result (before fetch) for next iteration + # Extract components from tuple result + current_alice_base = result[0] + current_bob_base = result[1] + current_alice_agg = result[2] + + print(f"\n✅ Training complete!") + print(f" Final loss: {loss_fetched}") + + # Summary + print("\n" + "=" * 80) + print("Key Takeaways") + print("=" * 80) + print( + "1. ✅ Vertical data partitioning: Alice has features + labels, Bob has features" + ) + print("2. ✅ Three-model architecture: Alice base, Bob base, Alice aggregate") + print( + "3. ✅ State dict pattern: model_state_dict + opt_state + step (following Tutorial 08)" + ) + print("4. ✅ Optax integration: Use optax.sgd() for optimizer state management") + print("5. ✅ Surrogate loss: Base models use dot product (grad · embedding)") + print("6. ✅ Privacy-preserving: Only embeddings and gradients are shared") + print("7. ✅ Complete training: Multiple iterations with decreasing loss") + print("8. ✅ Efficient: Pass results before fetch, only fetch for inspection") + print("\nSplit Learning enables collaborative training without sharing raw data!") + print("=" * 80) + + +if __name__ == "__main__": + main() From 1e2153b513530dd61edd7e677ef5ed02efe7f212 Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 09:05:22 +0000 Subject: [PATCH 02/10] reformat --- .../v1/device/08_initialize_nnx_on_device.py | 22 ++++++++-------- .../v1/device/09_split_learning_vertical.py | 26 +++++++++---------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index c6088a1d..2ad8b306 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -35,11 +35,11 @@ import jax import jax.numpy as jnp - -import mplang.v1 as mp import optax from flax import nnx +import mplang.v1 as mp + class SimpleMLP(nnx.Module): """Simple Multi-Layer Perceptron for demonstration.""" @@ -96,7 +96,7 @@ def _init(): ) # Split and convert state to pure dict - graphdef, state = nnx.split(model) + _graphdef, state = nnx.split(model) state_dict = state.to_pure_dict() print(f"[Device P0] Initialized model with {len(state_dict)} parameter groups") @@ -151,18 +151,18 @@ def _infer(x, params_dict): # So we can use it directly with replace_by_pure_dict abs_state.replace_by_pure_dict(params_dict) - print(f"[Device P0] Reconstructed state from dict") + print("[Device P0] Reconstructed state from dict") # Merge to get working model with the passed parameters model = nnx.merge(graphdef, abs_state) - print(f"[Device P0] Reconstructed model from state dict") - print(f"[Device P0] Running inference...") + print("[Device P0] Reconstructed model from state dict") + print("[Device P0] Running inference...") # Run inference output = model(x) - print(f"[Device P0] Inference complete!") + print("[Device P0] Inference complete!") print(f" Input shape: {x.shape}") print(f" Output shape: {output.shape}") @@ -191,7 +191,7 @@ def _init(): ) # Split and convert model state to pure dict - graphdef, state = nnx.split(model) + _graphdef, state = nnx.split(model) model_state_dict = state.to_pure_dict() # Initialize optimizer and convert its state to pure dict @@ -199,7 +199,7 @@ def _init(): tx = optax.sgd(learning_rate) opt_state = tx.init(model_state_dict) - print(f"[Device P0] Initialized model with optimizer") + print("[Device P0] Initialized model with optimizer") print(f"[Device P0] Model state keys: {list(model_state_dict.keys())}") print(f"[Device P0] Optimizer state type: {type(opt_state)}") @@ -340,7 +340,7 @@ def main(): if not isinstance(output, jnp.ndarray): output = jnp.array(output) - print(f"\n[Driver] Inference complete!") + print("\n[Driver] Inference complete!") print(f" Output shape: {output.shape}") print(f" Output sample: {output[0]}") @@ -393,7 +393,7 @@ def main(): # Fetch only for human inspection updated_train_dict, loss_value = mp.fetch(simulator, train_result) - print(f"\n[Driver] Training step complete!") + print("\n[Driver] Training step complete!") print(f" Loss value: {loss_value}") print( f" Updated model state keys: {list(updated_train_dict['model_state_dict'].keys())}" diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index dffcf6e4..7b58eda9 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -60,12 +60,12 @@ import jax import jax.numpy as jnp - -import mplang.v1 as mp import optax import pandas as pd from flax import nnx -from mplang.v1.core.dtypes import FLOAT64, INT64 + +import mplang.v1 as mp +from mplang.v1.core.dtypes import FLOAT64 from mplang.v1.ops import basic as basic_ops # We'll use a simplified state management approach similar to Tutorial 08 @@ -570,9 +570,9 @@ def prepare_vertical_split_data(): df_bob = pd.DataFrame(bob_cols) df_bob.to_csv(f"{base_dir}/bob_data.csv", index=False) - print(f"✅ Alice data: {N_SAMPLES} samples × {M1} features + labels") + print(f"✅ Alice data: {N_SAMPLES} samples x {M1} features + labels") print(f" Saved to: {base_dir}/alice_data.csv") - print(f"✅ Bob data: {N_SAMPLES} samples × {M2} features") + print(f"✅ Bob data: {N_SAMPLES} samples x {M2} features") print(f" Saved to: {base_dir}/bob_data.csv") print("=" * 80) @@ -638,7 +638,7 @@ def _init(): model = AliceBaseModel(input_dim=m1, hidden_dim=h1, rngs=nnx.Rngs(seed)) # Split into graphdef + state - graphdef, state = nnx.split(model) + _graphdef, state = nnx.split(model) # Convert state to pure dict model_state_dict = state.to_pure_dict() @@ -661,7 +661,7 @@ def _init(): model = BobBaseModel(input_dim=m2, hidden_dim=h2, rngs=nnx.Rngs(seed)) # Split into graphdef + state - graphdef, state = nnx.split(model) + _graphdef, state = nnx.split(model) # Convert state to pure dict model_state_dict = state.to_pure_dict() @@ -691,7 +691,7 @@ def _init(): ) # Split into graphdef + state - graphdef, state = nnx.split(model) + _graphdef, state = nnx.split(model) # Convert state to pure dict model_state_dict = state.to_pure_dict() @@ -889,7 +889,7 @@ def main(): SEED_AGG, LEARNING_RATE, ) - print(f"✅ Alice aggregate model initialized on P0 ({H1+H2} → 16 → {N_CLASSES})") + print(f"✅ Alice aggregate model initialized on P0 ({H1 + H2} → 16 → {N_CLASSES})") # Step 5: Run training iterations print("\n[Step 5] Running split learning training iterations...") @@ -931,9 +931,9 @@ def main(): # Fetch only for inspection ( - new_alice_base_fetched, - new_bob_base_fetched, - new_alice_agg_fetched, + _new_alice_base_fetched, + _new_bob_base_fetched, + _new_alice_agg_fetched, loss_fetched, ) = mp.fetch(simulator, result) @@ -945,7 +945,7 @@ def main(): current_bob_base = result[1] current_alice_agg = result[2] - print(f"\n✅ Training complete!") + print("\n✅ Training complete!") print(f" Final loss: {loss_fetched}") # Summary From e648e0548f35a1bcca6e1bcc09226a2a70a5282f Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 09:13:27 +0000 Subject: [PATCH 03/10] fixed based on comments --- .../v1/device/08_initialize_nnx_on_device.py | 41 +++---- .../v1/device/09_split_learning_vertical.py | 102 +++++------------- 2 files changed, 51 insertions(+), 92 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index 2ad8b306..26926a06 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -58,22 +58,20 @@ def __call__(self, x: jax.Array) -> jax.Array: # Cluster configuration -cluster_spec = mp.ClusterSpec.from_dict( - { - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, - }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, +cluster_spec = mp.ClusterSpec.from_dict({ + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, }, - } -) + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + }, +}) @mp.function @@ -369,9 +367,16 @@ def main(): # This gives the model something real to learn x_train = jax.random.normal(jax.random.PRNGKey(888), (32, 8)) # Create target with a simple linear relationship - true_weights = jnp.array( - [[2.0], [-1.0], [0.5], [1.5], [-0.5], [1.0], [-1.5], [0.0]] - ) + true_weights = jnp.array([ + [2.0], + [-1.0], + [0.5], + [1.5], + [-0.5], + [1.0], + [-1.5], + [0.0], + ]) true_weights2 = jnp.array([[1.0], [-0.5]]) y_train = ( x_train @ true_weights @ true_weights2.T diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index 7b58eda9..c9b8fa31 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -89,23 +89,21 @@ SEED_AGG = 44 # Cluster specification -cluster_spec = mp.ClusterSpec.from_dict( - { - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, - }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, - "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, +cluster_spec = mp.ClusterSpec.from_dict({ + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, }, - } -) + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, + }, +}) # ============================================================================ @@ -288,40 +286,6 @@ def bob_base_forward(x, model_dict, m2, h2, seed): return h2 -def alice_aggregate_forward_and_loss(h1, h2, y, model_dict, n_classes, seed): - """Alice's aggregate model forward pass with loss computation. - - Args: - h1: Alice's embeddings (n, h1) - h2: Bob's embeddings (n, h2) - y: Labels (n,) - model_dict: Alice aggregate model state as dict - n_classes: Number of output classes - seed: Random seed for model reconstruction - - Returns: - loss: Scalar loss value - logits: Model predictions (n, n_classes) - combined: Concatenated embeddings (n, h1+h2) - """ - # Reconstruct model - graphdef, state, _, _ = reconstruct_model_from_dict( - model_dict, AliceAggregateModel, H1 + H2, 16, n_classes, rngs=nnx.Rngs(seed) - ) - model = nnx.merge(graphdef, state) - - # Concatenate embeddings - combined = jnp.concatenate([h1, h2], axis=1) - - # Forward pass - logits = model(combined) - - # Compute loss (cross-entropy) - loss = jnp.mean(optax.softmax_cross_entropy_with_integer_labels(logits, y)) - - return loss, logits, combined - - # ============================================================================ # Helper Functions: Backward Pass # ============================================================================ @@ -365,12 +329,10 @@ def loss_fn(state_dict, h1, h2, y): # Get model state dict model_state_dict = model_dict["model_state_dict"] - # Compute gradients w.r.t. state_dict and inputs (h1, h2) - grad_fn = jax.grad(loss_fn, argnums=(0, 1, 2)) - grads_state_dict, grad_h1, grad_h2 = grad_fn(model_state_dict, h1, h2, y) - - # Also compute loss for logging - loss = loss_fn(model_state_dict, h1, h2, y) + # Compute loss and gradients w.r.t. state_dict and inputs (h1, h2) in a single pass + # Using value_and_grad is more efficient than computing loss twice + grad_fn = jax.value_and_grad(loss_fn, argnums=(0, 1, 2)) + loss, (grads_state_dict, grad_h1, grad_h2) = grad_fn(model_state_dict, h1, h2, y) return grads_state_dict, grad_h1, grad_h2, loss @@ -597,9 +559,10 @@ def load_vertical_split_data( """ # Define schemas (all columns must have same dtype for table_to_tensor) # Cast label to FLOAT64 here, convert back to int after splitting - schema_alice = mp.TableType.from_dict( - {**{f"alice_f{i}": FLOAT64 for i in range(m1)}, "label": FLOAT64} - ) + schema_alice = mp.TableType.from_dict({ + **{f"alice_f{i}": FLOAT64 for i in range(m1)}, + "label": FLOAT64, + }) schema_bob = mp.TableType.from_dict({f"bob_f{i}": FLOAT64 for i in range(m2)}) # Read CSVs as tables on respective devices @@ -640,14 +603,11 @@ def _init(): # Split into graphdef + state _graphdef, state = nnx.split(model) - # Convert state to pure dict - model_state_dict = state.to_pure_dict() - # Initialize optimizer tx = optax.sgd(learning_rate) - opt_state = tx.init(model_state_dict) + opt_state = tx.init(state.to_pure_dict()) - return {"model_state_dict": model_state_dict, "opt_state": opt_state, "step": 0} + return model_state_to_dict(state, opt_state, 0) return mp.device("P0", fe_type="nnx")(_init)() @@ -663,14 +623,11 @@ def _init(): # Split into graphdef + state _graphdef, state = nnx.split(model) - # Convert state to pure dict - model_state_dict = state.to_pure_dict() - # Initialize optimizer tx = optax.sgd(learning_rate) - opt_state = tx.init(model_state_dict) + opt_state = tx.init(state.to_pure_dict()) - return {"model_state_dict": model_state_dict, "opt_state": opt_state, "step": 0} + return model_state_to_dict(state, opt_state, 0) return mp.device("P1", fe_type="nnx")(_init)() @@ -693,14 +650,11 @@ def _init(): # Split into graphdef + state _graphdef, state = nnx.split(model) - # Convert state to pure dict - model_state_dict = state.to_pure_dict() - # Initialize optimizer tx = optax.sgd(learning_rate) - opt_state = tx.init(model_state_dict) + opt_state = tx.init(state.to_pure_dict()) - return {"model_state_dict": model_state_dict, "opt_state": opt_state, "step": 0} + return model_state_to_dict(state, opt_state, 0) return mp.device("P0", fe_type="nnx")(_init)() From dffb23ed99b0bfd6d9b592d0ee5d7c50ec72974d Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 09:22:29 +0000 Subject: [PATCH 04/10] fix all reported bugs --- .../v1/device/09_split_learning_vertical.py | 62 ++++++++++--------- 1 file changed, 33 insertions(+), 29 deletions(-) diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index c9b8fa31..27a598f8 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -89,21 +89,23 @@ SEED_AGG = 44 # Cluster specification -cluster_spec = mp.ClusterSpec.from_dict({ - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, +cluster_spec = mp.ClusterSpec.from_dict( + { + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, + }, + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, - "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, - }, -}) + } +) # ============================================================================ @@ -247,7 +249,7 @@ def alice_base_forward(x, model_dict, m1, h1, seed): seed: Random seed for model reconstruction Returns: - h1: Alice's embeddings (n, h1) + embeddings: Alice's embeddings (n, h1) """ # Reconstruct model graphdef, state, _, _ = reconstruct_model_from_dict( @@ -256,9 +258,9 @@ def alice_base_forward(x, model_dict, m1, h1, seed): model = nnx.merge(graphdef, state) # Forward pass - h1 = model(x) + embeddings = model(x) - return h1 + return embeddings def bob_base_forward(x, model_dict, m2, h2, seed): @@ -272,7 +274,7 @@ def bob_base_forward(x, model_dict, m2, h2, seed): seed: Random seed for model reconstruction Returns: - h2: Bob's embeddings (n, h2) + embeddings: Bob's embeddings (n, h2) """ # Reconstruct model graphdef, state, _, _ = reconstruct_model_from_dict( @@ -281,9 +283,9 @@ def bob_base_forward(x, model_dict, m2, h2, seed): model = nnx.merge(graphdef, state) # Forward pass - h2 = model(x) + embeddings = model(x) - return h2 + return embeddings # ============================================================================ @@ -305,7 +307,7 @@ def alice_aggregate_backward(h1, h2, y, model_dict, n_classes, seed): seed: Random seed for model reconstruction Returns: - grads_state: Gradients for aggregate model parameters + grads_state_dict: Gradients for aggregate model parameters grad_h1: Gradient w.r.t. Alice's embeddings (n, h1) grad_h2: Gradient w.r.t. Bob's embeddings (n, h2) loss: Scalar loss value @@ -559,10 +561,12 @@ def load_vertical_split_data( """ # Define schemas (all columns must have same dtype for table_to_tensor) # Cast label to FLOAT64 here, convert back to int after splitting - schema_alice = mp.TableType.from_dict({ - **{f"alice_f{i}": FLOAT64 for i in range(m1)}, - "label": FLOAT64, - }) + schema_alice = mp.TableType.from_dict( + { + **{f"alice_f{i}": FLOAT64 for i in range(m1)}, + "label": FLOAT64, + } + ) schema_bob = mp.TableType.from_dict({f"bob_f{i}": FLOAT64 for i in range(m2)}) # Read CSVs as tables on respective devices @@ -710,24 +714,24 @@ def split_learning_train_step( # === Forward Pass === # 1. Alice base model forward on P0 - h1 = mp.device("P0", fe_type="nnx")(alice_base_forward)( + h1_embeddings = mp.device("P0", fe_type="nnx")(alice_base_forward)( alice_features, alice_base_model_dict, m1, h1, seed_alice ) # 2. Bob base model forward on P1 - h2_on_p1 = mp.device("P1", fe_type="nnx")(bob_base_forward)( + h2_embeddings = mp.device("P1", fe_type="nnx")(bob_base_forward)( bob_features, bob_base_model_dict, m2, h2, seed_bob ) # 3. Transfer Bob's embeddings to Alice (P1 → P0) - h2_on_p0 = mp.put("P0", h2_on_p1) + h2_on_p0 = mp.put("P0", h2_embeddings) # === Backward Pass (Alice Aggregate Model) === # 4. Alice computes gradients for aggregate model using actual loss agg_grads_dict, grad_h1, grad_h2, loss = mp.device("P0", fe_type="nnx")( get_alice_agg_backward_fn() - )(h1, h2_on_p0, alice_labels, alice_agg_model_dict) + )(h1_embeddings, h2_on_p0, alice_labels, alice_agg_model_dict) # 5. Alice updates her aggregate model with actual gradients new_alice_agg_model_dict = mp.device("P0", fe_type="nnx")(update_model_state)( From b75e5fdfccdd9154296f7cb566ffff7220326345 Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 09:24:22 +0000 Subject: [PATCH 05/10] reformat --- .../v1/device/09_split_learning_vertical.py | 40 +++++++++---------- 1 file changed, 18 insertions(+), 22 deletions(-) diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index 27a598f8..f25032fc 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -89,23 +89,21 @@ SEED_AGG = 44 # Cluster specification -cluster_spec = mp.ClusterSpec.from_dict( - { - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, - }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, - "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, +cluster_spec = mp.ClusterSpec.from_dict({ + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, }, - } -) + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + "P1": {"kind": "PPU", "members": ["node_1"], "config": {}}, + }, +}) # ============================================================================ @@ -561,12 +559,10 @@ def load_vertical_split_data( """ # Define schemas (all columns must have same dtype for table_to_tensor) # Cast label to FLOAT64 here, convert back to int after splitting - schema_alice = mp.TableType.from_dict( - { - **{f"alice_f{i}": FLOAT64 for i in range(m1)}, - "label": FLOAT64, - } - ) + schema_alice = mp.TableType.from_dict({ + **{f"alice_f{i}": FLOAT64 for i in range(m1)}, + "label": FLOAT64, + }) schema_bob = mp.TableType.from_dict({f"bob_f{i}": FLOAT64 for i in range(m2)}) # Read CSVs as tables on respective devices From 61a2f81ce588eb515b5ec551d5cfb6f7c60fa05e Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 10:09:53 +0000 Subject: [PATCH 06/10] update structure --- .../v1/device/08_initialize_nnx_on_device.py | 610 ++++++++++-------- .../v1/device/09_split_learning_vertical.py | 405 +++++------- 2 files changed, 495 insertions(+), 520 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index 26926a06..c53bd5b4 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -40,6 +40,10 @@ import mplang.v1 as mp +# ============================================================================ +# Section 1: Pure Python Functions (Model Definition & Logic) +# ============================================================================ + class SimpleMLP(nnx.Module): """Simple Multi-Layer Perceptron for demonstration.""" @@ -57,82 +61,180 @@ def __call__(self, x: jax.Array) -> jax.Array: return x -# Cluster configuration -cluster_spec = mp.ClusterSpec.from_dict({ - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, - }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, - }, -}) +def init_model_logic(input_dim: int, hidden_dim: int, output_dim: int, seed: int): + """Pure function: Initialize model and return state as dict. + This is the core logic without MPLang decoration. -@mp.function -def initialize_model_and_return_state_dict( - input_dim: int, hidden_dim: int, output_dim: int, seed: int -): - """Initialize model on device P0 and return state as a pure dict. + Args: + input_dim: Model input dimension + hidden_dim: Model hidden dimension + output_dim: Model output dimension + seed: Random seed for initialization Returns: - Pure Python dict containing model parameters (can cross @mp.function boundaries) + Pure Python dict containing model parameters """ + # Create model + model = SimpleMLP( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + + # Split and convert state to pure dict + _graphdef, state = nnx.split(model) + state_dict = state.to_pure_dict() + + print(f"[Device P0] Initialized model with {len(state_dict)} parameter groups") + print(f"[Device P0] State dict keys: {list(state_dict.keys())}") + + # Return the state dict - it's in format {layer: {param: [array, None]}} + return state_dict + + +def inference_logic( + x: jax.Array, + params_dict: dict, + input_dim: int, + hidden_dim: int, + output_dim: int, + seed: int, +): + """Pure function: Run inference using state dict. + + This is the core logic without MPLang decoration. + + Args: + x: Input data [batch_size, input_dim] + params_dict: Pure Python dict containing model parameters + input_dim: Model input dimension + hidden_dim: Model hidden dimension + output_dim: Model output dimension + seed: Random seed (for GraphDef reconstruction) - def _init(): - # Create model on device - model = SimpleMLP( + Returns: + Model output logits + """ + # Create an abstract model to get the GraphDef (memory efficient!) + # This only creates the structure without allocating actual arrays + abs_model = nnx.eval_shape( + lambda: SimpleMLP( input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, rngs=nnx.Rngs(seed), ) + ) + graphdef, abs_state = nnx.split(abs_model) - # Split and convert state to pure dict - _graphdef, state = nnx.split(model) - state_dict = state.to_pure_dict() + print(f"[Device P0] Received state dict with keys: {list(params_dict.keys())}") - print(f"[Device P0] Initialized model with {len(state_dict)} parameter groups") - print(f"[Device P0] State dict keys: {list(state_dict.keys())}") + # The params_dict is already in the format from to_pure_dict(): {key: [array, None]} + # So we can use it directly with replace_by_pure_dict + abs_state.replace_by_pure_dict(params_dict) - # Return the state dict - it's in format {layer: {param: [array, None]}} - return state_dict + print("[Device P0] Reconstructed state from dict") - return mp.device("P0", fe_type="nnx")(_init)() + # Merge to get working model with the passed parameters + model = nnx.merge(graphdef, abs_state) + print("[Device P0] Reconstructed model from state dict") + print("[Device P0] Running inference...") -@mp.function -def run_inference_with_state_dict( - test_input: jax.Array, - state_dict: dict, + # Run inference + output = model(x) + + print("[Device P0] Inference complete!") + print(f" Input shape: {x.shape}") + print(f" Output shape: {output.shape}") + + return output + + +def init_model_with_optimizer_logic( + input_dim: int, hidden_dim: int, output_dim: int, seed: int, learning_rate: float +): + """Pure function: Initialize model + optimizer and return states as dict. + + This is the core logic without MPLang decoration. + + Args: + input_dim: Model input dimension + hidden_dim: Model hidden dimension + output_dim: Model output dimension + seed: Random seed for initialization + learning_rate: Learning rate for optimizer + + Returns: + Dict with model_state_dict, opt_state, and step + """ + # Create model + model = SimpleMLP( + input_dim=input_dim, + hidden_dim=hidden_dim, + output_dim=output_dim, + rngs=nnx.Rngs(seed), + ) + + # Split and convert model state to pure dict + _graphdef, state = nnx.split(model) + model_state_dict = state.to_pure_dict() + + # Initialize optimizer and convert its state to pure dict + # Note: optax states are already pytrees, but we'll store them explicitly + tx = optax.sgd(learning_rate) + opt_state = tx.init(model_state_dict) + + print("[Device P0] Initialized model with optimizer") + print(f"[Device P0] Model state keys: {list(model_state_dict.keys())}") + print(f"[Device P0] Optimizer state type: {type(opt_state)}") + + return { + "model_state_dict": model_state_dict, + "opt_state": opt_state, + "step": 0, + } + + +def train_step_logic( + train_state: dict, + batch_x: jax.Array, + batch_y: jax.Array, input_dim: int, hidden_dim: int, output_dim: int, seed: int, + learning_rate: float, ): - """Run inference using a state dict passed from outside. + """Pure function: Perform one training step. + + This is the core logic without MPLang decoration. + Forward → loss → backward → update Args: - test_input: Input data [batch_size, input_dim] - state_dict: Pure Python dict containing model parameters + train_state: Dict with model_state_dict, opt_state, step + batch_x: Input batch (batch_size, input_dim) + batch_y: Target batch (batch_size, output_dim) input_dim: Model input dimension hidden_dim: Model hidden dimension output_dim: Model output dimension seed: Random seed (for GraphDef reconstruction) + learning_rate: Learning rate for optimizer Returns: - Model output logits + Tuple of (updated_train_dict, loss_value) """ - - def _infer(x, params_dict): - """Inner function that takes explicit inputs.""" - # Create an abstract model to get the GraphDef (memory efficient!) - # This only creates the structure without allocating actual arrays + # Extract state + model_state_dict = train_state["model_state_dict"] + opt_state = train_state["opt_state"] + step = train_state["step"] + + # Define loss function that works with jax.grad + def loss_fn(state_dict, x, y): + # Reconstruct model from state dict for forward pass + # Create abstract model to get GraphDef (memory efficient!) abs_model = nnx.eval_shape( lambda: SimpleMLP( input_dim=input_dim, @@ -142,154 +244,162 @@ def _infer(x, params_dict): ) ) graphdef, abs_state = nnx.split(abs_model) + abs_state.replace_by_pure_dict(state_dict) + model = nnx.merge(graphdef, abs_state) - print(f"[Device P0] Received state dict with keys: {list(params_dict.keys())}") + # Forward pass + logits = model(x) - # The params_dict is already in the format from to_pure_dict(): {key: [array, None]} - # So we can use it directly with replace_by_pure_dict - abs_state.replace_by_pure_dict(params_dict) + # MSE loss + loss = jnp.mean((logits - y) ** 2) + return loss - print("[Device P0] Reconstructed state from dict") + # Compute loss and gradients in a single pass (efficient!) + loss_and_grad_fn = jax.value_and_grad(loss_fn) + loss_value, grads_dict = loss_and_grad_fn(model_state_dict, batch_x, batch_y) - # Merge to get working model with the passed parameters - model = nnx.merge(graphdef, abs_state) + # Apply optimizer update + tx = optax.sgd(learning_rate) + updates, new_opt_state = tx.update(grads_dict, opt_state, model_state_dict) + new_model_state_dict = optax.apply_updates(model_state_dict, updates) - print("[Device P0] Reconstructed model from state dict") - print("[Device P0] Running inference...") + # Return updated state + new_train_dict = { + "model_state_dict": new_model_state_dict, + "opt_state": new_opt_state, + "step": step + 1, + } - # Run inference - output = model(x) + return new_train_dict, loss_value - print("[Device P0] Inference complete!") - print(f" Input shape: {x.shape}") - print(f" Output shape: {output.shape}") - return output +# ============================================================================ +# Section 2: MPLang Functions (Multi-Party Device Execution) +# ============================================================================ - return mp.device("P0", fe_type="nnx")(_infer)(test_input, state_dict) +# Cluster configuration +cluster_spec = mp.ClusterSpec.from_dict({ + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, + }, + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + }, +}) @mp.function -def initialize_model_with_optimizer( - input_dim: int, hidden_dim: int, output_dim: int, seed: int, learning_rate: float +def demo_basic_model_init_and_inference( + input_dim: int, + hidden_dim: int, + output_dim: int, + seed: int, + test_input: jax.Array, ): - """Initialize model + optimizer on device P0 and return both states as dict. + """Combined MPLang function: Initialize model, run inference, return both. - Returns: - Dict with model_state_dict and opt_state (both as pure Python dicts) - """ + This combines two operations into one MPLang function to reduce overhead: + 1. Initialize model and get state dict + 2. Run inference with the state dict - def _init(): - # Create model on device - model = SimpleMLP( - input_dim=input_dim, - hidden_dim=hidden_dim, - output_dim=output_dim, - rngs=nnx.Rngs(seed), - ) + Args: + input_dim: Model input dimension + hidden_dim: Model hidden dimension + output_dim: Model output dimension + seed: Random seed for initialization + test_input: Input data for inference - # Split and convert model state to pure dict - _graphdef, state = nnx.split(model) - model_state_dict = state.to_pure_dict() + Returns: + Tuple of (state_dict, inference_output) + """ - # Initialize optimizer and convert its state to pure dict - # Note: optax states are already pytrees, but we'll store them explicitly - tx = optax.sgd(learning_rate) - opt_state = tx.init(model_state_dict) + def _combined(): + # Step 1: Initialize model + state_dict = init_model_logic(input_dim, hidden_dim, output_dim, seed) - print("[Device P0] Initialized model with optimizer") - print(f"[Device P0] Model state keys: {list(model_state_dict.keys())}") - print(f"[Device P0] Optimizer state type: {type(opt_state)}") + # Step 2: Run inference using the same state dict + output = inference_logic( + test_input, state_dict, input_dim, hidden_dim, output_dim, seed + ) - return { - "model_state_dict": model_state_dict, - "opt_state": opt_state, - "step": 0, - } + return state_dict, output - return mp.device("P0", fe_type="nnx")(_init)() + return mp.device("P0", fe_type="nnx")(_combined)() @mp.function -def train_step( - train_dict: dict, - x: jnp.ndarray, - y: jnp.ndarray, +def train_model_for_n_steps( input_dim: int, hidden_dim: int, output_dim: int, seed: int, learning_rate: float, + x_train: jax.Array, + y_train: jax.Array, + n_steps: int, ): - """Perform one training step: forward → loss → backward → update. + """Combined MPLang function: Initialize model with optimizer and train for N steps. + + This combines the entire training pipeline into one MPLang function: + 1. Initialize model with optimizer + 2. Run N training steps in a loop + 3. Return final state and loss history + + This is much more efficient than calling train_step N times across @mp.function boundaries! Args: - train_dict: Dict with model_state_dict, opt_state, step - x: Input batch (batch_size, input_dim) - y: Target batch (batch_size, output_dim) input_dim: Model input dimension hidden_dim: Model hidden dimension output_dim: Model output dimension - seed: Random seed (for GraphDef reconstruction) + seed: Random seed for initialization learning_rate: Learning rate for optimizer + x_train: Training input data + y_train: Training target data + n_steps: Number of training steps to run Returns: - Tuple of (updated_train_dict, loss_value) + Tuple of (final_train_dict, loss_history) """ - def _train(train_state, batch_x, batch_y): - # Extract state - model_state_dict = train_state["model_state_dict"] - opt_state = train_state["opt_state"] - step = train_state["step"] - - # Define loss function that works with jax.grad - def loss_fn(state_dict, x, y): - # Reconstruct model from state dict for forward pass - # Create abstract model to get GraphDef (memory efficient!) - abs_model = nnx.eval_shape( - lambda: SimpleMLP( - input_dim=input_dim, - hidden_dim=hidden_dim, - output_dim=output_dim, - rngs=nnx.Rngs(seed), - ) - ) - graphdef, abs_state = nnx.split(abs_model) - abs_state.replace_by_pure_dict(state_dict) - model = nnx.merge(graphdef, abs_state) - - # Forward pass - logits = model(x) - - # MSE loss - loss = jnp.mean((logits - y) ** 2) - return loss + def _train_loop(): + # Initialize model with optimizer + train_dict = init_model_with_optimizer_logic( + input_dim, hidden_dim, output_dim, seed, learning_rate + ) - # Compute gradients - loss_value = loss_fn(model_state_dict, batch_x, batch_y) - grad_fn = jax.grad(loss_fn) - grads_dict = grad_fn(model_state_dict, batch_x, batch_y) + # Training loop - all steps run on device without crossing boundaries + loss_history = [] + for _step_idx in range(n_steps): + train_dict, loss_value = train_step_logic( + train_dict, + x_train, + y_train, + input_dim, + hidden_dim, + output_dim, + seed, + learning_rate, + ) + loss_history.append(loss_value) - # Apply optimizer update - tx = optax.sgd(learning_rate) - updates, new_opt_state = tx.update(grads_dict, opt_state, model_state_dict) - new_model_state_dict = optax.apply_updates(model_state_dict, updates) + return train_dict, jnp.array(loss_history) - # Return updated state - new_train_dict = { - "model_state_dict": new_model_state_dict, - "opt_state": new_opt_state, - "step": step + 1, - } + return mp.device("P0", fe_type="nnx")(_train_loop)() - return new_train_dict, loss_value - return mp.device("P0", fe_type="nnx")(_train)(train_dict, x, y) +# ============================================================================ +# Section 3: Main Execution Flow +# ============================================================================ def main(): - """Main demonstration: initialize model on P0, then run inference using state dict.""" + """Main demonstration: Shows both efficient batched operations and granular operations.""" print("=" * 80) print("NNX Model with State Dict Pattern") print("=" * 80) @@ -297,76 +407,51 @@ def main(): # Setup simulator simulator = mp.Simulator(cluster_spec) - # Step 1: Initialize model on device P0 and get state as dict - print("\n[Step 1] Initializing model on device P0...") - state_dict_result = mp.evaluate( - simulator, - initialize_model_and_return_state_dict, - input_dim=8, - hidden_dim=16, - output_dim=2, - seed=42, - ) - - # Fetch the state dict to driver - state_dict = mp.fetch(simulator, state_dict_result) - print(f"\n[Driver] Received state dict with keys: {list(state_dict.keys())}") + # ========================================================================= + # PART A: Efficient Batched Operations (RECOMMENDED) + # ========================================================================= + print("\n" + "=" * 80) + print("PART A: Efficient Batched MPLang Functions (RECOMMENDED)") + print("=" * 80) + print("Combining multiple operations into single @mp.function calls") + print("reduces overhead and improves performance!\n") - # Step 2: Run inference using the state dict - print("\n[Step 2] Running inference with state dict...") + # Step A1: Combined init + inference in ONE MPLang call + print("[Step A1] Initialize model + Run inference (combined operation)...") # Create test input test_key = jax.random.PRNGKey(999) test_input = jax.random.normal(test_key, (4, 8)) - output_result = mp.evaluate( + combined_result = mp.evaluate( simulator, - run_inference_with_state_dict, - test_input, - state_dict_result, + demo_basic_model_init_and_inference, input_dim=8, hidden_dim=16, output_dim=2, - seed=42, # Same seed to get same GraphDef + seed=42, + test_input=test_input, ) - output = mp.fetch(simulator, output_result) + state_dict, output = mp.fetch(simulator, combined_result) - # Handle output if it's wrapped + # Handle wrapped output if isinstance(output, (list, tuple)) and len(output) > 0: output = output[0] if not isinstance(output, jnp.ndarray): output = jnp.array(output) - print("\n[Driver] Inference complete!") - print(f" Output shape: {output.shape}") - print(f" Output sample: {output[0]}") - - # Step 3: Initialize model with optimizer - print("\n[Step 3] Initializing model with optimizer on device P0...") - train_state_dict_result = mp.evaluate( - simulator, - initialize_model_with_optimizer, - input_dim=8, - hidden_dim=16, - output_dim=2, - seed=42, - learning_rate=0.01, - ) - - # Fetch the train state dict to driver - train_state_dict = mp.fetch(simulator, train_state_dict_result) - print( - f"\n[Driver] Received train state dict with keys: {list(train_state_dict.keys())}" - ) + print("✅ Model initialized and inference complete!") + print(f" State dict keys: {list(state_dict.keys())}") + print(f" Output shape: {output.shape}") + print(f" Output sample: {output[0]}") - # Step 4: Perform a single training step - print("\n[Step 4] Performing a single training step...") + # Step A2: Combined training loop in ONE MPLang call + print("\n[Step A2] Initialize + Train for 10 steps (combined operation)...") + print("This is MUCH more efficient than 10 separate @mp.function calls!\n") - # Create a simple learnable dataset: y = 2*x + noise - # This gives the model something real to learn + # Create training data x_train = jax.random.normal(jax.random.PRNGKey(888), (32, 8)) - # Create target with a simple linear relationship true_weights = jnp.array([ [2.0], [-1.0], @@ -385,105 +470,70 @@ def main(): train_result = mp.evaluate( simulator, - train_step, - train_state_dict_result, - x_train, - y_train, + train_model_for_n_steps, input_dim=8, hidden_dim=16, output_dim=2, seed=42, learning_rate=0.01, + x_train=x_train, + y_train=y_train, + n_steps=10, ) - # Fetch only for human inspection - updated_train_dict, loss_value = mp.fetch(simulator, train_result) - print("\n[Driver] Training step complete!") - print(f" Loss value: {loss_value}") - print( - f" Updated model state keys: {list(updated_train_dict['model_state_dict'].keys())}" - ) - print(f" Updated optimizer state type: {type(updated_train_dict['opt_state'])}") - print(f" Step number: {updated_train_dict['step']}") - - # Step 5: Demonstrate Optimizer State Persists Across Multiple Training Steps - print("\n" + "=" * 80) - print("[Step 5] Optimizer State Persists Across MPLang Function Boundaries") - print("=" * 80) - print("Demonstrating that optimizer state can be passed back and forth") - print("\nRunning 5 more training steps with THE SAME batch...") - print("Note: Loss should decrease steadily as the model learns the pattern!\n") - - # Keep using train_result (before fetch) for subsequent computations - current_train_result = train_result - - for step_idx in range(5): - # THE KEY POINT: Pass train_result[0] (the train_dict) to next mp.evaluate! - # train_result is a tuple (train_dict, loss), we need just the train_dict - # Use the SAME training data to show actual learning - current_train_result = mp.evaluate( - simulator, - train_step, - current_train_result[ - 0 - ], # Pass just the train_dict (first element of tuple) - x_train, # Same data - so we can see the model actually learning! - y_train, - input_dim=8, - hidden_dim=16, - output_dim=2, - seed=42, - learning_rate=0.01, - ) + final_train_dict, loss_history = mp.fetch(simulator, train_result) - # Fetch only for display (human inspection) - updated_train_dict, loss_value = mp.fetch(simulator, current_train_result) - print( - f" Step {step_idx + 2}: Loss = {loss_value}, Global step = {updated_train_dict['step']}" - ) + # Handle wrapped arrays + if isinstance(loss_history, (list, tuple)): + loss_history = loss_history[0] if len(loss_history) > 0 else loss_history + if not isinstance(loss_history, jnp.ndarray): + loss_history = jnp.array(loss_history) - # Final fetch for summary - final_train_dict, final_loss = mp.fetch(simulator, current_train_result) + print("\n✅ Training complete!") + print(f" Final step: {final_train_dict['step']}") + print(f" Initial loss: {loss_history[0]:.4f}") + print(f" Final loss: {loss_history[-1]:.4f}") + print(f" Loss reduction: {(loss_history[0] - loss_history[-1]):.4f}") + print( + f" All losses: {[f'{loss:.4f}' for loss in loss_history[:5]]}... (showing first 5)" + ) + # Summary print("\n" + "=" * 80) - print("Training Progress Summary") + print("Key Takeaways") print("=" * 80) - print(f"Final loss: {final_loss}") - print(f"Final optimizer state type: {type(final_train_dict['opt_state'])}") - print(f"Final global step: {final_train_dict['step']}") + print("\n✅ NNX + MPLANG PATTERNS:") + print("1. ✅ Separate pure Python logic from MPLang decorators") + print(" - Section 1: Pure functions (testable, reusable)") + print(" - Section 2: MPLang wrappers (device placement)") + print(" - Section 3: Main execution (orchestration)") + + print("\n2. ✅ Batch operations into larger @mp.function calls") + print(" - demo_basic_model_init_and_inference: init + inference together") + print(" - train_model_for_n_steps: init + N training steps together") + print(" - Reduces @mp.function boundary crossings = better performance!") + + print("\n3. ✅ Use state dicts for cross-boundary communication") + print(" - state.to_pure_dict() → serialize model state") + print(" - state.replace_by_pure_dict() → deserialize model state") + print(" - Works with optimizer state too!") + + print("\n4. ✅ Use nnx.eval_shape() for memory efficiency") + print(" - Creates GraphDef without allocating arrays") + print(" - Especially important when reconstructing models repeatedly") + + print("\n5. ✅ Use jax.value_and_grad for efficient gradient computation") + print(" - Computes loss and gradients in a single pass") + print(" - More efficient than separate forward and backward passes") + + print("\n🎯 BEST PRACTICES:") + print(" • Batch related operations into single @mp.function calls") + print(" • Minimize @mp.function boundary crossings") + print(" • Group operations together for efficiency") + print(" • Use state dicts for flexible model management") + print(" • Leverage nnx.eval_shape() for memory-efficient model reconstruction") print("\n" + "=" * 80) - print("Key Takeaways") - print("=" * 80) - print("1. ✅ Initialize NNX models on devices with mp.device('P0', fe_type='nnx')") - print("2. ✅ Convert state to pure dict: state.to_pure_dict()") - print(" - Format: {layer: {param: [array, None]}}") - print("3. ✅ Pass dict across @mp.function boundaries (works in MPLang v1!)") - print("4. ✅ Reconstruct state: state.replace_by_pure_dict(state_dict)") - print("5. ✅ Use nnx.eval_shape() for memory-efficient GraphDef creation") - print(" - Creates abstract model without allocating arrays!") - print("6. ✅ Merge with GraphDef: model = nnx.merge(graphdef, state)") - print("7. ✅ This enables efficient model reuse without reinitializing weights!") - print("8. ✅ Initialize model with optimizer state management") - print(" - Use optax for flexible optimizer integration") - print(" - Manage optimizer state as pure Python dicts") - print( - "9. ✅ Perform training steps with full control: forward, loss, backward, update" - ) - print(" - Stateful training across @mp.function boundaries") - print(" - Checkpoint-style updates with pure dicts") - print( - "10. ✅ CRITICAL: Optimizer state persists across MPLang function boundaries!" - ) - print(" - Pass train_dict (with opt_state) between training steps") - print(" - Optimizer accumulates state (momentum, etc.) correctly") - print(" - Step counter tracks global training progress") - print(" - Enables iterative training with stateful optimizers") - print(" - Loss decreases steadily when training on same batch!") - print("\nThis pattern solves the MPLang v1 limitation of passing complex types") - print("by using pure Python dicts that can cross @mp.function boundaries.") - print("Optimizer state management enables proper multi-step training!") - print("=" * 80) if __name__ == "__main__": diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index f25032fc..8c9d6c73 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -188,8 +188,6 @@ def __call__(self, x: jax.Array) -> jax.Array: def model_state_to_dict(state, opt_state, step): """Convert model state to pure Python dict for cross-device transfer. - Note: graphdef is NOT stored - it can be reconstructed from model class - Args: state: NNX State (model parameters) opt_state: Optax optimizer state @@ -197,6 +195,10 @@ def model_state_to_dict(state, opt_state, step): Returns: Dict with keys: model_state_dict, opt_state, step + + Note: graphdef is NOT included - it can be reconstructed from the model class + definition. Only the state dict (parameters), optimizer state, and step counter + are stored for transfer. """ return { "model_state_dict": state.to_pure_dict(), @@ -296,6 +298,10 @@ def alice_aggregate_backward(h1, h2, y, model_dict, n_classes, seed): This is standard supervised learning - compute loss from predictions and labels. + Note: This function captures module-level constant H1 and H2 from the outer scope + for model reconstruction. The n_classes and seed parameters are passed explicitly + to support testing and reuse in different contexts. + Args: h1: Alice's embeddings (n, h1) h2: Bob's embeddings (n, h2) @@ -314,7 +320,8 @@ def alice_aggregate_backward(h1, h2, y, model_dict, n_classes, seed): # Define loss function def loss_fn(state_dict, h1, h2, y): # Reconstruct model from state dict - # n_classes and seed are captured from outer scope and are concrete values + # Note: H1, H2 are captured from module scope as concrete values + # n_classes and seed are captured from function parameters temp_model = AliceAggregateModel(H1 + H2, 16, n_classes, rngs=nnx.Rngs(seed)) graphdef, temp_state = nnx.split(temp_model) temp_state.replace_by_pure_dict(state_dict) @@ -361,7 +368,8 @@ def alice_base_backward(x_alice, grad_h1, model_dict, m1, h1, seed): # Surrogate loss function def surrogate_loss_fn(state_dict, x, grad_from_next_layer): # Reconstruct model from state dict - # m1, h1, seed are captured from outer scope and are concrete values + # m1, h1, seed are function parameters captured by this closure. + # They are fixed values for this invocation (not JAX tracers). temp_model = AliceBaseModel(m1, h1, rngs=nnx.Rngs(seed)) graphdef, temp_state = nnx.split(temp_model) temp_state.replace_by_pure_dict(state_dict) @@ -406,7 +414,8 @@ def bob_base_backward(x_bob, grad_h2, model_dict, m2, h2, seed): # Surrogate loss function def surrogate_loss_fn(state_dict, x, grad_from_next_layer): # Reconstruct model from state dict - # m2, h2, seed are captured from outer scope and are concrete values + # m2, h2, seed are function parameters captured by this closure. + # They are fixed values for this invocation (not JAX tracers). temp_model = BobBaseModel(m2, h2, rngs=nnx.Rngs(seed)) graphdef, temp_state = nnx.split(temp_model) temp_state.replace_by_pure_dict(state_dict) @@ -558,7 +567,8 @@ def load_vertical_split_data( bob_features: Tensor on P1 (n, m2) """ # Define schemas (all columns must have same dtype for table_to_tensor) - # Cast label to FLOAT64 here, convert back to int after splitting + # Convert label to FLOAT64 to match feature columns, then cast back to int after + # This workaround is needed because table_to_tensor requires uniform dtypes schema_alice = mp.TableType.from_dict({ **{f"alice_f{i}": FLOAT64 for i in range(m1)}, "label": FLOAT64, @@ -588,185 +598,125 @@ def split_features_labels(data): # ============================================================================ -# Model Initialization +# Batched MPLang Functions (Efficient Training) # ============================================================================ @mp.function -def initialize_alice_base_model(m1: int, h1: int, seed: int, learning_rate: float): - """Initialize Alice's base model on P0 and return state as dict.""" +def train_split_learning_for_n_steps( + alice_features, + alice_labels, + bob_features, + m1: int, + m2: int, + h1: int, + h2: int, + n_classes: int, + seed_alice: int, + seed_bob: int, + seed_agg: int, + learning_rate: float, + n_steps: int, +): + """Batched training: Initialize all models + train for N steps in one MPLang call. - def _init(): - # Create model - model = AliceBaseModel(input_dim=m1, hidden_dim=h1, rngs=nnx.Rngs(seed)) + This is the most efficient approach, combining: + 1. Initialize all three models + 2. Run N training iterations + All in a single @mp.function call, avoiding repeated boundary crossings. - # Split into graphdef + state - _graphdef, state = nnx.split(model) + Returns: + Tuple of (final_alice_base_dict, final_bob_base_dict, final_alice_agg_dict, final_loss) + """ + + # === Step 1: Initialize all models === - # Initialize optimizer + def _init_alice_base(): + model = AliceBaseModel(input_dim=m1, hidden_dim=h1, rngs=nnx.Rngs(seed_alice)) + _graphdef, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) - return model_state_to_dict(state, opt_state, 0) - return mp.device("P0", fe_type="nnx")(_init)() - - -@mp.function -def initialize_bob_base_model(m2: int, h2: int, seed: int, learning_rate: float): - """Initialize Bob's base model on P1 and return state as dict.""" - - def _init(): - # Create model - model = BobBaseModel(input_dim=m2, hidden_dim=h2, rngs=nnx.Rngs(seed)) - - # Split into graphdef + state + def _init_bob_base(): + model = BobBaseModel(input_dim=m2, hidden_dim=h2, rngs=nnx.Rngs(seed_bob)) _graphdef, state = nnx.split(model) - - # Initialize optimizer tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) - return model_state_to_dict(state, opt_state, 0) - return mp.device("P1", fe_type="nnx")(_init)() - - -@mp.function -def initialize_alice_agg_model( - input_dim: int, hidden_dim: int, output_dim: int, seed: int, learning_rate: float -): - """Initialize Alice's aggregate model on P0 and return state as dict.""" - - def _init(): - # Create model + def _init_alice_agg(): model = AliceAggregateModel( - input_dim=input_dim, - hidden_dim=hidden_dim, - output_dim=output_dim, - rngs=nnx.Rngs(seed), + input_dim=h1 + h2, + hidden_dim=16, + output_dim=n_classes, + rngs=nnx.Rngs(seed_agg), ) - - # Split into graphdef + state _graphdef, state = nnx.split(model) - - # Initialize optimizer tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) - return model_state_to_dict(state, opt_state, 0) - return mp.device("P0", fe_type="nnx")(_init)() - - -# ============================================================================ -# Training Step -# ============================================================================ - - -@mp.function -def split_learning_train_step( - alice_features, # On P0: (10000, m1) - alice_labels, # On P0: (10000,) - bob_features, # On P1: (10000, m2) - alice_base_model_dict, # Alice base model state as dict - bob_base_model_dict, # Bob base model state as dict - alice_agg_model_dict, # Alice aggregate model state as dict - learning_rate: float, - m1: int, - m2: int, - h1: int, - h2: int, - n_classes: int, - seed_alice: int, - seed_bob: int, - seed_agg: int, -): - """One complete split learning training step (single batch = full dataset). - - Training Flow (Split Learning): - 1. Alice computes base model forward → h1 - 2. Bob computes base model forward → h2 - 3. Transfer h2 from P1 to P0 - 4. Alice runs aggregate model → logits → loss (actual supervised loss) - 5. Alice computes gradients for aggregate model (∂L/∂params_agg, ∂L/∂h1, ∂L/∂h2) - 6. Alice updates aggregate model with actual gradients - 7. Alice computes gradients for her base model using surrogate loss (grad_h1 · h1) - 8. Alice updates her base model with surrogate gradients - 9. Alice sends ∂L/∂h2 back to Bob (P0 → P1) - 10. Bob computes gradients for his base model using surrogate loss (grad_h2 · h2) - 11. Bob updates his base model with surrogate gradients - - Note: Uses helper functions (get_*_backward_fn) that bind module-level constants - via functools.partial to avoid traced parameter issues in JAX gradient computation. - - Returns: - new_alice_base_model_dict: Updated Alice base model state (as dict) - new_bob_base_model_dict: Updated Bob base model state (as dict) - new_alice_agg_model_dict: Updated Alice aggregate model state (as dict) - loss: Training loss (scalar) - """ + # Initialize models + alice_base_dict = mp.device("P0", fe_type="nnx")(_init_alice_base)() + alice_agg_dict = mp.device("P0", fe_type="nnx")(_init_alice_agg)() + bob_base_dict = mp.device("P1", fe_type="nnx")(_init_bob_base)() - # === Forward Pass === + # === Step 2: Training loop (all iterations run inside this @mp.function) === - # 1. Alice base model forward on P0 - h1_embeddings = mp.device("P0", fe_type="nnx")(alice_base_forward)( - alice_features, alice_base_model_dict, m1, h1, seed_alice - ) + final_loss = None - # 2. Bob base model forward on P1 - h2_embeddings = mp.device("P1", fe_type="nnx")(bob_base_forward)( - bob_features, bob_base_model_dict, m2, h2, seed_bob - ) - - # 3. Transfer Bob's embeddings to Alice (P1 → P0) - h2_on_p0 = mp.put("P0", h2_embeddings) + for _step_idx in range(n_steps): + # Forward pass: Alice base + h1_embeddings = mp.device("P0", fe_type="nnx")(alice_base_forward)( + alice_features, alice_base_dict, m1, h1, seed_alice + ) - # === Backward Pass (Alice Aggregate Model) === + # Forward pass: Bob base + h2_embeddings = mp.device("P1", fe_type="nnx")(bob_base_forward)( + bob_features, bob_base_dict, m2, h2, seed_bob + ) - # 4. Alice computes gradients for aggregate model using actual loss - agg_grads_dict, grad_h1, grad_h2, loss = mp.device("P0", fe_type="nnx")( - get_alice_agg_backward_fn() - )(h1_embeddings, h2_on_p0, alice_labels, alice_agg_model_dict) + # Transfer h2: P1 → P0 + h2_on_p0 = mp.put("P0", h2_embeddings) - # 5. Alice updates her aggregate model with actual gradients - new_alice_agg_model_dict = mp.device("P0", fe_type="nnx")(update_model_state)( - alice_agg_model_dict, agg_grads_dict, learning_rate - ) + # Backward pass: Alice aggregate model (actual supervised loss) + agg_grads_dict, grad_h1, grad_h2, loss = mp.device("P0", fe_type="nnx")( + get_alice_agg_backward_fn() + )(h1_embeddings, h2_on_p0, alice_labels, alice_agg_dict) - # === Backward Pass (Alice Base Model) === + # Update: Alice aggregate model + alice_agg_dict = mp.device("P0", fe_type="nnx")(update_model_state)( + alice_agg_dict, agg_grads_dict, learning_rate + ) - # 6. Alice computes gradients for her base model using surrogate loss - alice_base_grads_dict = mp.device("P0", fe_type="nnx")( - get_alice_base_backward_fn() - )(alice_features, grad_h1, alice_base_model_dict) + # Backward pass: Alice base model (surrogate loss) + alice_base_grads_dict = mp.device("P0", fe_type="nnx")( + get_alice_base_backward_fn() + )(alice_features, grad_h1, alice_base_dict) - # 7. Alice updates her base model with surrogate gradients - new_alice_base_model_dict = mp.device("P0", fe_type="nnx")(update_model_state)( - alice_base_model_dict, alice_base_grads_dict, learning_rate - ) + # Update: Alice base model + alice_base_dict = mp.device("P0", fe_type="nnx")(update_model_state)( + alice_base_dict, alice_base_grads_dict, learning_rate + ) - # === Backward Pass (Bob Base Model) === + # Transfer grad_h2: P0 → P1 + grad_h2_on_p1 = mp.put("P1", grad_h2) - # 8. Transfer grad_h2 to Bob (P0 → P1) - grad_h2_on_p1 = mp.put("P1", grad_h2) + # Backward pass: Bob base model (surrogate loss) + bob_base_grads_dict = mp.device("P1", fe_type="nnx")( + get_bob_base_backward_fn() + )(bob_features, grad_h2_on_p1, bob_base_dict) - # 9. Bob computes gradients for his base model using surrogate loss - bob_base_grads_dict = mp.device("P1", fe_type="nnx")(get_bob_base_backward_fn())( - bob_features, grad_h2_on_p1, bob_base_model_dict - ) + # Update: Bob base model + bob_base_dict = mp.device("P1", fe_type="nnx")(update_model_state)( + bob_base_dict, bob_base_grads_dict, learning_rate + ) - # 10. Bob updates his base model with surrogate gradients - new_bob_base_model_dict = mp.device("P1", fe_type="nnx")(update_model_state)( - bob_base_model_dict, bob_base_grads_dict, learning_rate - ) + # Track final loss + final_loss = loss - return ( - new_alice_base_model_dict, - new_bob_base_model_dict, - new_alice_agg_model_dict, - loss, - ) + return alice_base_dict, bob_base_dict, alice_agg_dict, final_loss # ============================================================================ @@ -808,116 +758,91 @@ def main(): ) print("✅ Data loaded successfully") - # Step 4: Initialize models - print("\n[Step 4] Initializing models on devices...") + # Step 4: Train with batched approach (efficient) + print("\n[Step 4] Initialize all models + Train for 10 steps (batched)...") + print("This runs ALL operations in a single @mp.function call for efficiency!\n") - # Alice base model (P0) - alice_base_model_dict = mp.evaluate( + batched_result = mp.evaluate( simulator, - initialize_alice_base_model, + train_split_learning_for_n_steps, + alice_features, + alice_labels, + bob_features, M1, - H1, - SEED_ALICE, - LEARNING_RATE, - ) - print(f"✅ Alice base model initialized on P0 (m1={M1} → h1={H1})") - - # Bob base model (P1) - bob_base_model_dict = mp.evaluate( - simulator, - initialize_bob_base_model, M2, + H1, H2, - SEED_BOB, - LEARNING_RATE, - ) - print(f"✅ Bob base model initialized on P1 (m2={M2} → h2={H2})") - - # Alice aggregate model (P0) - alice_agg_model_dict = mp.evaluate( - simulator, - initialize_alice_agg_model, - H1 + H2, # Concatenated embeddings - 16, # Hidden dimension N_CLASSES, + SEED_ALICE, + SEED_BOB, SEED_AGG, LEARNING_RATE, + n_steps=10, ) - print(f"✅ Alice aggregate model initialized on P0 ({H1 + H2} → 16 → {N_CLASSES})") - - # Step 5: Run training iterations - print("\n[Step 5] Running split learning training iterations...") - print("Training flow:") - print(" 1. Alice forward: x_alice → h1") - print(" 2. Bob forward: x_bob → h2") - print(" 3. Transfer h2: P1 → P0") - print(" 4. Alice aggregate: [h1, h2] → loss") - print(" 5. Alice aggregate backward & update") - print(" 6. Alice base backward & update (surrogate loss)") - print(" 7. Transfer grad_h2: P0 → P1") - print(" 8. Bob base backward & update (surrogate loss)") - print("\nRunning 5 training iterations on the same batch...") - - current_alice_base = alice_base_model_dict - current_bob_base = bob_base_model_dict - current_alice_agg = alice_agg_model_dict - - for iter_idx in range(5): - result = mp.evaluate( - simulator, - split_learning_train_step, - alice_features, - alice_labels, - bob_features, - current_alice_base, - current_bob_base, - current_alice_agg, - LEARNING_RATE, - M1, - M2, - H1, - H2, - N_CLASSES, - SEED_ALICE, - SEED_BOB, - SEED_AGG, - ) - # Fetch only for inspection - ( - _new_alice_base_fetched, - _new_bob_base_fetched, - _new_alice_agg_fetched, - loss_fetched, - ) = mp.fetch(simulator, result) + ( + final_alice_base, + final_bob_base, + final_alice_agg, + final_loss, + ) = mp.fetch(simulator, batched_result) - print(f" Iteration {iter_idx + 1}: Loss = {loss_fetched}") - - # Update current state using result (before fetch) for next iteration - # Extract components from tuple result - current_alice_base = result[0] - current_bob_base = result[1] - current_alice_agg = result[2] + # Handle final_loss if it's wrapped + if isinstance(final_loss, (list, tuple)): + final_loss = final_loss[0] if len(final_loss) > 0 else 0.0 print("\n✅ Training complete!") - print(f" Final loss: {loss_fetched}") + print( + f" Final steps: Alice={final_alice_base['step']}, " + f"Bob={final_bob_base['step']}, " + f"Agg={final_alice_agg['step']}" + ) + print(f" Final loss: {final_loss:.4f}") # Summary print("\n" + "=" * 80) print("Key Takeaways") print("=" * 80) - print( - "1. ✅ Vertical data partitioning: Alice has features + labels, Bob has features" - ) - print("2. ✅ Three-model architecture: Alice base, Bob base, Alice aggregate") - print( - "3. ✅ State dict pattern: model_state_dict + opt_state + step (following Tutorial 08)" - ) - print("4. ✅ Optax integration: Use optax.sgd() for optimizer state management") - print("5. ✅ Surrogate loss: Base models use dot product (grad · embedding)") - print("6. ✅ Privacy-preserving: Only embeddings and gradients are shared") - print("7. ✅ Complete training: Multiple iterations with decreasing loss") - print("8. ✅ Efficient: Pass results before fetch, only fetch for inspection") + print("\n✅ SPLIT LEARNING CONCEPTS:") + print("1. ✅ Vertical data partitioning:") + print(" - Alice has features + labels, Bob has features") + print(" - Privacy-preserving: only embeddings and gradients are shared") + + print("\n2. ✅ Three-model architecture:") + print(" - Alice base: m1 → h1 embeddings") + print(" - Bob base: m2 → h2 embeddings") + print(" - Alice aggregate: [h1, h2] → logits (has labels, computes loss)") + + print("\n3. ✅ State dict pattern (following Tutorial 08):") + print(" - model_state_dict + opt_state + step") + print(" - Use state.to_pure_dict() and state.replace_by_pure_dict()") + print(" - Enables passing state across @mp.function boundaries") + + print("\n4. ✅ Surrogate loss for base models:") + print(" - Alice base: L_surrogate = grad_h1 · h1") + print(" - Bob base: L_surrogate = grad_h2 · h2") + print(" - Mimics backpropagation from aggregate model") + + print("\n✅ MPLANG EFFICIENCY PATTERNS:") + print("5. ✅ Batch operations into large @mp.function calls:") + print(" - train_split_learning_for_n_steps: init + N training steps in one call") + print(" - Minimizes @mp.function boundary crossings") + print(" - Significantly reduces overhead for multi-step training") + + print("\n6. ✅ Use nnx.eval_shape() for memory efficiency:") + print(" - Creates GraphDef without allocating arrays") + print(" - Important when reconstructing models repeatedly") + + print("\n7. ✅ Use module-level constants + functools.partial:") + print(" - Avoids JAX tracer issues in gradient computation") + print(" - Bind constants at module level, not as traced parameters") + + print("\n🎯 BEST PRACTICES:") + print(" • Batch related operations into single @mp.function calls") + print(" • Minimize cross-device communication and boundary crossings") + print(" • Use state dicts for flexible model state management") + print(" • Leverage surrogate loss for privacy-preserving training") + print("\nSplit Learning enables collaborative training without sharing raw data!") print("=" * 80) From f4202dd3a709542cd0b4a26c7867ccd5367ce970 Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Thu, 18 Dec 2025 10:13:44 +0000 Subject: [PATCH 07/10] update docs --- .../v1/device/08_initialize_nnx_on_device.py | 68 ++++++++++--------- 1 file changed, 36 insertions(+), 32 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index c53bd5b4..f426752c 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -35,10 +35,10 @@ import jax import jax.numpy as jnp -import optax -from flax import nnx import mplang.v1 as mp +import optax +from flax import nnx # ============================================================================ # Section 1: Pure Python Functions (Model Definition & Logic) @@ -278,20 +278,22 @@ def loss_fn(state_dict, x, y): # ============================================================================ # Cluster configuration -cluster_spec = mp.ClusterSpec.from_dict({ - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, +cluster_spec = mp.ClusterSpec.from_dict( + { + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, + }, + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, - }, -}) + } +) @mp.function @@ -399,7 +401,7 @@ def _train_loop(): def main(): - """Main demonstration: Shows both efficient batched operations and granular operations.""" + """Main demonstration: Shows efficient batched operations for NNX models.""" print("=" * 80) print("NNX Model with State Dict Pattern") print("=" * 80) @@ -408,16 +410,16 @@ def main(): simulator = mp.Simulator(cluster_spec) # ========================================================================= - # PART A: Efficient Batched Operations (RECOMMENDED) + # Efficient Batched Operations (RECOMMENDED) # ========================================================================= print("\n" + "=" * 80) - print("PART A: Efficient Batched MPLang Functions (RECOMMENDED)") + print("Efficient Batched MPLang Functions (RECOMMENDED)") print("=" * 80) print("Combining multiple operations into single @mp.function calls") print("reduces overhead and improves performance!\n") - # Step A1: Combined init + inference in ONE MPLang call - print("[Step A1] Initialize model + Run inference (combined operation)...") + # Step 1: Combined init + inference in ONE MPLang call + print("[Step 1] Initialize model + Run inference (combined operation)...") # Create test input test_key = jax.random.PRNGKey(999) @@ -446,22 +448,24 @@ def main(): print(f" Output shape: {output.shape}") print(f" Output sample: {output[0]}") - # Step A2: Combined training loop in ONE MPLang call - print("\n[Step A2] Initialize + Train for 10 steps (combined operation)...") + # Step 2: Combined training loop in ONE MPLang call + print("\n[Step 2] Initialize + Train for 10 steps (combined operation)...") print("This is MUCH more efficient than 10 separate @mp.function calls!\n") # Create training data x_train = jax.random.normal(jax.random.PRNGKey(888), (32, 8)) - true_weights = jnp.array([ - [2.0], - [-1.0], - [0.5], - [1.5], - [-0.5], - [1.0], - [-1.5], - [0.0], - ]) + true_weights = jnp.array( + [ + [2.0], + [-1.0], + [0.5], + [1.5], + [-0.5], + [1.0], + [-1.5], + [0.0], + ] + ) true_weights2 = jnp.array([[1.0], [-0.5]]) y_train = ( x_train @ true_weights @ true_weights2.T From ad3afbc3e7377c5c99b7c7b3b22553c8f0aa69cb Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Dec 2025 10:33:14 +0800 Subject: [PATCH 08/10] Address review feedback: clarify documentation and fix type consistency (#306) * Initial plan * Address review comments: improve documentation and clarity Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com> * Fix type consistency for final_loss initialization Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com> --- .../v1/device/08_initialize_nnx_on_device.py | 15 ++--- .../v1/device/09_split_learning_vertical.py | 59 +++++++++++-------- 2 files changed, 42 insertions(+), 32 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index f426752c..2065042f 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -117,8 +117,9 @@ def inference_logic( Returns: Model output logits """ - # Create an abstract model to get the GraphDef (memory efficient!) - # This only creates the structure without allocating actual arrays + # Create an abstract model to get the GraphDef. + # nnx.eval_shape only creates the abstract structure; parameter arrays are + # materialized later when replace_by_pure_dict is called with params_dict. abs_model = nnx.eval_shape( lambda: SimpleMLP( input_dim=input_dim, @@ -233,8 +234,8 @@ def train_step_logic( # Define loss function that works with jax.grad def loss_fn(state_dict, x, y): - # Reconstruct model from state dict for forward pass - # Create abstract model to get GraphDef (memory efficient!) + # Reconstruct model from state dict for forward pass. + # Create abstract model to get GraphDef; parameters are then materialized from state_dict. abs_model = nnx.eval_shape( lambda: SimpleMLP( input_dim=input_dim, @@ -522,9 +523,9 @@ def main(): print(" - state.replace_by_pure_dict() → deserialize model state") print(" - Works with optimizer state too!") - print("\n4. ✅ Use nnx.eval_shape() for memory efficiency") - print(" - Creates GraphDef without allocating arrays") - print(" - Especially important when reconstructing models repeatedly") + print("\n4. ✅ Use nnx.eval_shape() for GraphDef reconstruction") + print(" - Creates abstract model structure for obtaining GraphDef") + print(" - Actual parameters are materialized when replace_by_pure_dict is called") print("\n5. ✅ Use jax.value_and_grad for efficient gradient computation") print(" - Computes loss and gradients in a single pass") diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index 8c9d6c73..3b1320d4 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -33,10 +33,9 @@ Key Concepts: 1. **State Management**: Use graphdef + state dict pattern (following Tutorial 08) - 2. **Memory Efficiency**: Use nnx.eval_shape() to create abstract models (no array allocation) - 3. **Surrogate Loss**: Base models use dot product (grad · embedding) for training - 4. **Privacy**: Only embeddings and gradients are shared, not raw features - 5. **Optimizer**: Recreated on each step (SGD with fixed learning rate) + 2. **Surrogate Loss**: Base models use dot product (grad · embedding) for training + 3. **Privacy**: Only embeddings and gradients are shared, not raw features + 4. **Optimizer**: Recreated on each step (SGD with fixed learning rate) Training Flow (One Iteration): 1. Alice forward: x_alice → h1 @@ -197,8 +196,10 @@ def model_state_to_dict(state, opt_state, step): Dict with keys: model_state_dict, opt_state, step Note: graphdef is NOT included - it can be reconstructed from the model class - definition. Only the state dict (parameters), optimizer state, and step counter - are stored for transfer. + definition, but requires the same model initialization parameters (input_dim, + hidden_dim, output_dim, seed) to be provided when reconstructing via + reconstruct_model_from_dict. Only the state dict (parameters), optimizer state, + and step counter are stored for transfer. """ return { "model_state_dict": state.to_pure_dict(), @@ -298,9 +299,10 @@ def alice_aggregate_backward(h1, h2, y, model_dict, n_classes, seed): This is standard supervised learning - compute loss from predictions and labels. - Note: This function captures module-level constant H1 and H2 from the outer scope + Note: This function captures module-level constants H1 and H2 from the outer scope for model reconstruction. The n_classes and seed parameters are passed explicitly - to support testing and reuse in different contexts. + (rather than captured from module scope) to support testing flexibility and avoid + closure issues with JAX tracing. Args: h1: Alice's embeddings (n, h1) @@ -319,9 +321,10 @@ def alice_aggregate_backward(h1, h2, y, model_dict, n_classes, seed): # Define loss function def loss_fn(state_dict, h1, h2, y): - # Reconstruct model from state dict - # Note: H1, H2 are captured from module scope as concrete values - # n_classes and seed are captured from function parameters + # Reconstruct model from state dict. + # Note: H1, H2 are module-level constants (safe to capture, not JAX tracers). + # n_classes and seed are parameters to alice_aggregate_backward and are + # captured here as fixed values for this invocation (avoiding tracer issues). temp_model = AliceAggregateModel(H1 + H2, 16, n_classes, rngs=nnx.Rngs(seed)) graphdef, temp_state = nnx.split(temp_model) temp_state.replace_by_pure_dict(state_dict) @@ -566,9 +569,18 @@ def load_vertical_split_data( alice_labels: Tensor on P0 (n,) bob_features: Tensor on P1 (n, m2) """ - # Define schemas (all columns must have same dtype for table_to_tensor) - # Convert label to FLOAT64 to match feature columns, then cast back to int after - # This workaround is needed because table_to_tensor requires uniform dtypes + # Define schemas (all columns must have same dtype for table_to_tensor). + # NOTE: mp.ops.basic.table_to_tensor requires that all columns in a table share + # the same dtype. Since Alice's feature columns are FLOAT64, we also declare + # the label column as FLOAT64 so that we can load features + labels in a single + # table_to_tensor call. In this tutorial, the labels in the CSV are integer-valued + # class IDs (e.g. 0/1), so converting the label column from FLOAT64 back to + # jnp.int32 after tensorization is an exact, lossless cast. + # + # If your labels are not integer-valued, or if they must be stored with a + # different dtype than the features, prefer a more robust pattern such as + # reading features and labels via separate schemas / read calls so you do not + # rely on a FLOAT64 → int cast. schema_alice = mp.TableType.from_dict({ **{f"alice_f{i}": FLOAT64 for i in range(m1)}, "label": FLOAT64, @@ -603,7 +615,7 @@ def split_features_labels(data): @mp.function -def train_split_learning_for_n_steps( +def initialize_and_train_split_learning( alice_features, alice_labels, bob_features, @@ -618,9 +630,10 @@ def train_split_learning_for_n_steps( learning_rate: float, n_steps: int, ): - """Batched training: Initialize all models + train for N steps in one MPLang call. + """Initialize all models and train for N steps in one batched MPLang call. - This is the most efficient approach, combining: + This is the most efficient approach, combining model initialization and training + into a single @mp.function call to minimize boundary crossings. 1. Initialize all three models 2. Run N training iterations All in a single @mp.function call, avoiding repeated boundary crossings. @@ -664,7 +677,7 @@ def _init_alice_agg(): # === Step 2: Training loop (all iterations run inside this @mp.function) === - final_loss = None + final_loss = jnp.array(0.0) # Initialize as JAX array to maintain consistent type for _step_idx in range(n_steps): # Forward pass: Alice base @@ -764,7 +777,7 @@ def main(): batched_result = mp.evaluate( simulator, - train_split_learning_for_n_steps, + initialize_and_train_split_learning, alice_features, alice_labels, bob_features, @@ -825,15 +838,11 @@ def main(): print("\n✅ MPLANG EFFICIENCY PATTERNS:") print("5. ✅ Batch operations into large @mp.function calls:") - print(" - train_split_learning_for_n_steps: init + N training steps in one call") + print(" - initialize_and_train_split_learning: init + N training steps in one call") print(" - Minimizes @mp.function boundary crossings") print(" - Significantly reduces overhead for multi-step training") - print("\n6. ✅ Use nnx.eval_shape() for memory efficiency:") - print(" - Creates GraphDef without allocating arrays") - print(" - Important when reconstructing models repeatedly") - - print("\n7. ✅ Use module-level constants + functools.partial:") + print("\n6. ✅ Use module-level constants + functools.partial:") print(" - Avoids JAX tracer issues in gradient computation") print(" - Bind constants at module level, not as traced parameters") From cd67a82f59b766b0a599eb14212598e1781fee22 Mon Sep 17 00:00:00 2001 From: "zoupeicheng.zpc" Date: Fri, 19 Dec 2025 02:35:28 +0000 Subject: [PATCH 09/10] ruff format --- .../v1/device/08_initialize_nnx_on_device.py | 54 +++++++++---------- .../v1/device/09_split_learning_vertical.py | 4 +- 2 files changed, 28 insertions(+), 30 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index 2065042f..3a4b76aa 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -35,11 +35,11 @@ import jax import jax.numpy as jnp - -import mplang.v1 as mp import optax from flax import nnx +import mplang.v1 as mp + # ============================================================================ # Section 1: Pure Python Functions (Model Definition & Logic) # ============================================================================ @@ -279,22 +279,20 @@ def loss_fn(state_dict, x, y): # ============================================================================ # Cluster configuration -cluster_spec = mp.ClusterSpec.from_dict( - { - "nodes": [ - {"name": "node_0", "endpoint": "127.0.0.1:61920"}, - {"name": "node_1", "endpoint": "127.0.0.1:61921"}, - ], - "devices": { - "SP0": { - "kind": "SPU", - "members": ["node_0", "node_1"], - "config": {"protocol": "SEMI2K", "field": "FM128"}, - }, - "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, +cluster_spec = mp.ClusterSpec.from_dict({ + "nodes": [ + {"name": "node_0", "endpoint": "127.0.0.1:61920"}, + {"name": "node_1", "endpoint": "127.0.0.1:61921"}, + ], + "devices": { + "SP0": { + "kind": "SPU", + "members": ["node_0", "node_1"], + "config": {"protocol": "SEMI2K", "field": "FM128"}, }, - } -) + "P0": {"kind": "PPU", "members": ["node_0"], "config": {}}, + }, +}) @mp.function @@ -455,18 +453,16 @@ def main(): # Create training data x_train = jax.random.normal(jax.random.PRNGKey(888), (32, 8)) - true_weights = jnp.array( - [ - [2.0], - [-1.0], - [0.5], - [1.5], - [-0.5], - [1.0], - [-1.5], - [0.0], - ] - ) + true_weights = jnp.array([ + [2.0], + [-1.0], + [0.5], + [1.5], + [-0.5], + [1.0], + [-1.5], + [0.0], + ]) true_weights2 = jnp.array([[1.0], [-0.5]]) y_train = ( x_train @ true_weights @ true_weights2.T diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index 3b1320d4..68d141ea 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -838,7 +838,9 @@ def main(): print("\n✅ MPLANG EFFICIENCY PATTERNS:") print("5. ✅ Batch operations into large @mp.function calls:") - print(" - initialize_and_train_split_learning: init + N training steps in one call") + print( + " - initialize_and_train_split_learning: init + N training steps in one call" + ) print(" - Minimizes @mp.function boundary crossings") print(" - Significantly reduces overhead for multi-step training") From 85cdb7a8e2f047ee4d42d5981f7cb2eda6427ba2 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Dec 2025 10:44:34 +0800 Subject: [PATCH 10/10] Replace _graphdef with _ for unused variables in NNX tutorials (#307) * Initial plan * Replace _graphdef with _ for unused variables Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com> --- tutorials/v1/device/08_initialize_nnx_on_device.py | 4 ++-- tutorials/v1/device/09_split_learning_vertical.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index 3a4b76aa..8fdd26fc 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -84,7 +84,7 @@ def init_model_logic(input_dim: int, hidden_dim: int, output_dim: int, seed: int ) # Split and convert state to pure dict - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) state_dict = state.to_pure_dict() print(f"[Device P0] Initialized model with {len(state_dict)} parameter groups") @@ -180,7 +180,7 @@ def init_model_with_optimizer_logic( ) # Split and convert model state to pure dict - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) model_state_dict = state.to_pure_dict() # Initialize optimizer and convert its state to pure dict diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index 68d141ea..e41252c8 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -646,14 +646,14 @@ def initialize_and_train_split_learning( def _init_alice_base(): model = AliceBaseModel(input_dim=m1, hidden_dim=h1, rngs=nnx.Rngs(seed_alice)) - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) return model_state_to_dict(state, opt_state, 0) def _init_bob_base(): model = BobBaseModel(input_dim=m2, hidden_dim=h2, rngs=nnx.Rngs(seed_bob)) - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) return model_state_to_dict(state, opt_state, 0) @@ -665,7 +665,7 @@ def _init_alice_agg(): output_dim=n_classes, rngs=nnx.Rngs(seed_agg), ) - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) return model_state_to_dict(state, opt_state, 0)