fix: setting it up like before

Author: init-22

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -13,6 +13,8 @@
 from algoperf.workloads.criteo1tb.criteo1tb_jax import models
 from algoperf.workloads.criteo1tb.workload import \
     BaseCriteo1TbDlrmSmallWorkload
+from custom_pytorch_jax_converter import use_pytorch_weights_inplace, use_pytorch_weights_inplace_mnist
+
 
 
 class Criteo1TbDlrmSmallWorkload(BaseCriteo1TbDlrmSmallWorkload):
@@ -103,8 +105,10 @@ def init_model_fn(
         {'params': params_rng, 'dropout': dropout_rng},
         jnp.ones(input_shape, jnp.float32))
     initial_params = initial_variables['params']
+    initial_params = use_pytorch_weights_inplace(initial_params, file_name="/results/pytorch_base_model_criteo1tb_24_june.pth")
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
+    return initial_params, None
     return jax_utils.replicate(initial_params), None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:

custom_pytorch_jax_converter.py

@@ -0,0 +1,186 @@
+import torch
+import numpy as np
+import jax
+import jax.numpy as jnp
+import logging
+import copy
+import copy
+from jax.tree_util import tree_map
+"""
+Jax default parameter structure:
+dict_keys(['Dense_0', 'Dense_1', 'Dense_2', 'Dense_3', 'Dense_4', 'Dense_5', 'Dense_6', 'Dense_7', 'embedding_table'])
+
+Pytorch stateduct structure:
+dict_keys(['embedding_chunk_0', 'embedding_chunk_1', 'embedding_chunk_2', 'embedding_chunk_3', 'bot_mlp.0.weight', 'bot_mlp.0.bias', 'bot_mlp.2.weight', 'bot_mlp.2.bias', 'bot_mlp.4.weight', 'bot_mlp.4.bias', 'top_mlp.0.weight', 'top_mlp.0.bias', 'top_mlp.2.weight', 'top_mlp.2.bias', 'top_mlp.4.weight', 'top_mlp.4.bias', 'top_mlp.6.weight', 'top_mlp.6.bias', 'top_mlp.8.weight', 'top_mlp.8.bias'])
+
+
+
+The following function converts the PyTorch weights to the Jax format
+and assigns them to the Jax model parameters.
+The function assumes that the Jax model parameters are already initialized
+and that the PyTorch weights are in the correct format.
+"""
+def use_pytorch_weights_inplace(jax_params, file_name=None, replicate=False):
+
+    # Load PyTorch state_dict
+    state_dict = torch.load(file_name)
+    print(state_dict.keys())
+    # Convert PyTorch tensors to NumPy arrays
+    numpy_weights = {k: v.cpu().numpy() for k, v in state_dict.items()}   
+
+    # --- Embedding Table ---
+    embedding_table = np.concatenate([
+        numpy_weights[f'embedding_chunk_{i}'] for i in range(4)
+    ], axis=0)  # adjust axis depending on chunking direction
+
+    jax_params['embedding_table'] = jnp.array(embedding_table)
+
+    # --- Bot MLP: Dense_0 to Dense_2 ---
+    for i, j in zip([0, 2, 4], range(3)):
+        jax_params[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'bot_mlp.{i}.weight'].T)
+        jax_params[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'bot_mlp.{i}.bias'])
+
+    # --- Top MLP: Dense_3 to Dense_7 ---
+    for i, j in zip([0, 2, 4, 6, 8], range(3, 8)):
+        jax_params[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'top_mlp.{i}.weight'].T)
+        jax_params[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'top_mlp.{i}.bias'])
+    #jax_params = tree_map(lambda x: jnp.array(x), jax_params) 
+    del state_dict
+    return jax_params
+
+
+def use_pytorch_weights_inplace_mnist(jax_params, file_name=None, replicate=False):
+    # Load the PyTorch checkpoint
+    ckpt = torch.load(file_name)
+    state_dict = ckpt['state_dict'] if 'state_dict' in ckpt else ckpt
+
+    print("Loaded PyTorch keys:", state_dict.keys())
+
+    # Convert to numpy
+    numpy_weights = {k: v.cpu().numpy() for k, v in state_dict.items()}
+
+    # Mapping PyTorch keys → JAX Dense layers
+    layer_map = {
+        'net.layer1': 'Dense_0',
+        'net.layer2': 'Dense_1',
+    }
+
+    for pt_name, jax_name in layer_map.items():
+        weight_key = f"{pt_name}.weight"
+        bias_key = f"{pt_name}.bias"
+
+        if weight_key not in numpy_weights or bias_key not in numpy_weights:
+            raise KeyError(f"Missing keys: {weight_key} or {bias_key} in PyTorch weights")
+
+        jax_params[jax_name]['kernel'] = jnp.array(numpy_weights[weight_key].T)  # Transpose!
+        jax_params[jax_name]['bias'] = jnp.array(numpy_weights[bias_key])
+
+    return jax_params
+
+
+# def are_weights_equal(params1, params2, atol=1e-6, rtol=1e-6):
+#     """Compares two JAX PyTrees of weights and prints where they differ."""
+#     all_equal = True
+
+#     def compare_fn(p1, p2):
+#         nonlocal all_equal
+#         #if not jnp.allclose(p1, p2):
+#         if not jnp.allclose(p1, p2, atol=atol, rtol=rtol):
+#             logging.info("❌ Mismatch found:")
+#             logging.info(f"Shape 1: {p1.shape}, Shape 2: {p2.shape}")
+#             logging.info(f"Max diff: {jnp.max(jnp.abs(p1 - p2))}")
+#             all_equal = False
+#         return jnp.allclose(p1, p2, atol=atol, rtol=rtol)
+
+#     try:
+#         _ = jax.tree_util.tree_map(compare_fn, params1, params2)
+#     except Exception as e:
+#         logging.info("❌ Structure mismatch or error during comparison:", e)
+#         return False
+
+#     if all_equal:
+#         logging.info("✅ All weights are equal (within tolerance)")
+#     return all_equal
+
+import jax
+import jax.numpy as jnp
+import logging
+
+def maybe_unreplicate(pytree):
+    """If leading axis matches device count, strip it assuming it's pmap replication."""
+    num_devices = jax.device_count()
+    return jax.tree_util.tree_map(
+        lambda x: x[0] if isinstance(x, jax.Array) and x.shape[0] == num_devices else x,
+        pytree
+    )
+
+def move_to_cpu(tree):
+    return jax.tree_util.tree_map(lambda x: jax.device_put(x, device=jax.devices("cpu")[0]), tree)
+
+
+def are_weights_equal(params1, params2, atol=1e-6, rtol=1e-6):
+    """Compares two JAX PyTrees of weights and logs where they differ, safely handling PMAP replication."""
+    # Attempt to unreplicate if needed
+    params1 = maybe_unreplicate(params1)
+    params2 = maybe_unreplicate(params2)
+
+    params1 = move_to_cpu(params1)
+    params2 = move_to_cpu(params2)
+
+    all_equal = True
+
+    def compare_fn(p1, p2):
+        nonlocal all_equal
+        if not jnp.allclose(p1, p2, atol=atol, rtol=rtol):
+            logging.info("❌ Mismatch found:")
+            logging.info(f"Shape 1: {p1.shape}, Shape 2: {p2.shape}")
+            logging.info(f"Max diff: {jnp.max(jnp.abs(p1 - p2))}")
+            all_equal = False
+        return jnp.allclose(p1, p2, atol=atol, rtol=rtol)
+
+    try:
+        jax.tree_util.tree_map(compare_fn, params1, params2)
+    except Exception as e:
+        logging.info("❌ Structure mismatch or error during comparison:", exc_info=True)
+        return False
+
+    if all_equal:
+        logging.info("✅ All weights are equal (within tolerance)")
+    return all_equal
+
+
+
+def use_pytorch_weights2(jax_params, file_name=None, replicate=False):
+
+    def deep_copy_to_cpu(pytree):
+        return tree_map(lambda x: jax.device_put(jnp.array(copy.deepcopy(x)), device=jax.devices("cpu")[0]), pytree)
+
+    breakpoint()
+    jax_copy = deep_copy_to_cpu(jax_params) 
+    # Load PyTorch state_dict lazily to CPU
+    state_dict = torch.load(file_name, map_location='cpu')
+    print(state_dict.keys())
+    # Convert PyTorch tensors to NumPy arrays
+    numpy_weights = {k: v.cpu().numpy() for k, v in state_dict.items()}   
+
+    # --- Embedding Table ---
+    embedding_table = np.concatenate([
+        numpy_weights[f'embedding_chunk_{i}'] for i in range(4)
+    ], axis=0)  # adjust axis depending on chunking direction
+
+    jax_copy['embedding_table'] = jnp.array(embedding_table)
+
+    # --- Bot MLP: Dense_0 to Dense_2 ---
+    for i, j in zip([0, 2, 4], range(3)):
+        jax_copy[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'bot_mlp.{i}.weight'].T)
+        jax_copy[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'bot_mlp.{i}.bias'])
+
+    # --- Top MLP: Dense_3 to Dense_7 ---
+    for i, j in zip([0, 2, 4, 6, 8], range(3, 8)):
+        jax_copy[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'top_mlp.{i}.weight'].T)
+        jax_copy[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'top_mlp.{i}.bias'])
+    #jax_copy = tree_map(lambda x: jnp.array(x), jax_copy) 
+    del state_dict
+
+    return jax_copy
+

reference_algorithms/schedule_free/jax/submission.py

@@ -10,6 +10,7 @@
 import jax.numpy as jnp
 from optax.contrib import schedule_free_adamw
 from algoperf import spec
+from custom_pytorch_jax_converter import use_pytorch_weights2, are_weights_equal
 
 _GRAD_CLIP_EPS = 1e-6
 
@@ -168,6 +169,15 @@ def update_params(workload: spec.Workload,
             'loss': loss[0],
             'grad_norm': grad_norm[0],
         }, global_step)
+    
+  # Log the number of parameters.
+  if global_step % 100 == 0 and workload.metrics_logger is not None:
+    date_ = "2025=06-14"
+    file_name = f"/results/schedule_free_test_pytorch_weights/criteo1tb_{date_}_after_{global_step}.pth"
+    params = use_pytorch_weights2(new_params, file_name=file_name, replicate=True)
+    are_weights_equal(new_params, params)
+    del params
+    
   return (new_optimizer_state, opt_update_fn), new_params, new_model_state
 
 

reference_algorithms/schedule_free/pytorch/submission.py

@@ -266,6 +266,14 @@ def closure():
 
   loss = optimizer_state['optimizer'].step(closure)
 
+  # # current_state  = current_model.module.state_dict() is a workaround for DDP
+  # if global_step==1:
+  #   torch.save(current_model.module.state_dict(), "/results/pytorch_base_model_criteo1tb_11may_global_step2.pth")
+  #   import torch.distributed as dist
+  #   import sys
+  #   dist.destroy_process_group()
+  #   sys.exit(0)
+
   return (optimizer_state, current_param_container, new_model_state)