Skip to content

Commit c2ca92a

Browse files
init-22init-22
committed
fix: changes
1 parent c057f10 commit c2ca92a

7 files changed

Lines changed: 262 additions & 154 deletions

File tree

‎algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from algoperf.workloads.criteo1tb.criteo1tb_jax import models
1414
from algoperf.workloads.criteo1tb.workload import \
1515
BaseCriteo1TbDlrmSmallWorkload
16-
from custom_pytorch_jax_converter import use_pytorch_weights
16+
from custom_pytorch_jax_converter import use_pytorch_weights_inplace
1717

1818

1919
class Criteo1TbDlrmSmallWorkload(BaseCriteo1TbDlrmSmallWorkload):
@@ -105,7 +105,8 @@ def init_model_fn(
105105
jnp.ones(input_shape, jnp.float32))
106106
initial_params = initial_variables['params']
107107
logging.info('\n\nInitializing with Pytorch weights\n\n')
108-
initial_params = use_pytorch_weights(initial_params, file_name="/results/pytorch_base_model_criteo1tb_8_june.pth")
108+
breakpoint()
109+
initial_params = use_pytorch_weights_inplace(initial_params, file_name="/results/pytorch_base_model_criteo1tb_13_june.pth")
109110
self._param_shapes = param_utils.jax_param_shapes(initial_params)
110111
self._param_types = param_utils.jax_param_types(self._param_shapes)
111112
return jax_utils.replicate(initial_params), None

‎algoperf/workloads/criteo1tb/criteo1tb_pytorch/workload.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ def init_model_fn(
8888
dropout_rate=dropout_rate,
8989
use_layer_norm=self.use_layer_norm,
9090
embedding_init_multiplier=self.embedding_init_multiplier)
91-
torch.save(model.state_dict(), "/results/pytorch_base_model_criteo1tb_8_june.pth")
91+
torch.save(model.state_dict(), "/results/pytorch_base_model_criteo1tb_13_june.pth")
9292
self._param_shapes = param_utils.pytorch_param_shapes(model)
9393
self._param_types = param_utils.pytorch_param_types(self._param_shapes)
9494
model.to(DEVICE)

‎algoperf/workloads/criteo1tb/workload.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@ def _build_input_queue(
115115
repeat_final_dataset: Optional[bool] = None,
116116
num_batches: Optional[int] = None) -> Iterator[Dict[str, spec.Tensor]]:
117117
del cache
118+
118119
ds = input_pipeline.get_criteo1tb_dataset(
119120
split=split,
120121
shuffle_rng=data_rng,

‎custom_pytorch_jax_converter.py‎

Lines changed: 104 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22
import numpy as np
33
import jax
44
import jax.numpy as jnp
5+
import logging
6+
import copy
7+
import copy
8+
from jax.tree_util import tree_map
59
"""
610
Jax default parameter structure:
711
dict_keys(['Dense_0', 'Dense_1', 'Dense_2', 'Dense_3', 'Dense_4', 'Dense_5', 'Dense_6', 'Dense_7', 'embedding_table'])
@@ -16,7 +20,8 @@
1620
The function assumes that the Jax model parameters are already initialized
1721
and that the PyTorch weights are in the correct format.
1822
"""
19-
def use_pytorch_weights(jax_params, file_name=None):
23+
def use_pytorch_weights_inplace(jax_params, file_name=None, replicate=False):
24+
2025
# Load PyTorch state_dict
2126
state_dict = torch.load(file_name)
2227
print(state_dict.keys())
@@ -28,41 +33,125 @@ def use_pytorch_weights(jax_params, file_name=None):
2833
numpy_weights[f'embedding_chunk_{i}'] for i in range(4)
2934
], axis=0) # adjust axis depending on chunking direction
3035

31-
jax_params['embedding_table'] = embedding_table
36+
jax_params['embedding_table'] = jnp.array(embedding_table)
3237

3338
# --- Bot MLP: Dense_0 to Dense_2 ---
3439
for i, j in zip([0, 2, 4], range(3)):
35-
jax_params[f'Dense_{j}']['kernel'] = numpy_weights[f'bot_mlp.{i}.weight'].T
36-
jax_params[f'Dense_{j}']['bias'] = numpy_weights[f'bot_mlp.{i}.bias']
40+
jax_params[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'bot_mlp.{i}.weight'].T)
41+
jax_params[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'bot_mlp.{i}.bias'])
3742

3843
# --- Top MLP: Dense_3 to Dense_7 ---
3944
for i, j in zip([0, 2, 4, 6, 8], range(3, 8)):
40-
jax_params[f'Dense_{j}']['kernel'] = numpy_weights[f'top_mlp.{i}.weight'].T
41-
jax_params[f'Dense_{j}']['bias'] = numpy_weights[f'top_mlp.{i}.bias']
42-
45+
jax_params[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'top_mlp.{i}.weight'].T)
46+
jax_params[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'top_mlp.{i}.bias'])
47+
#jax_params = tree_map(lambda x: jnp.array(x), jax_params)
48+
del state_dict
4349
return jax_params
4450

4551

52+
# def are_weights_equal(params1, params2, atol=1e-6, rtol=1e-6):
53+
# """Compares two JAX PyTrees of weights and prints where they differ."""
54+
# all_equal = True
55+
56+
# def compare_fn(p1, p2):
57+
# nonlocal all_equal
58+
# #if not jnp.allclose(p1, p2):
59+
# if not jnp.allclose(p1, p2, atol=atol, rtol=rtol):
60+
# logging.info("❌ Mismatch found:")
61+
# logging.info(f"Shape 1: {p1.shape}, Shape 2: {p2.shape}")
62+
# logging.info(f"Max diff: {jnp.max(jnp.abs(p1 - p2))}")
63+
# all_equal = False
64+
# return jnp.allclose(p1, p2, atol=atol, rtol=rtol)
65+
66+
# try:
67+
# _ = jax.tree_util.tree_map(compare_fn, params1, params2)
68+
# except Exception as e:
69+
# logging.info("❌ Structure mismatch or error during comparison:", e)
70+
# return False
71+
72+
# if all_equal:
73+
# logging.info("✅ All weights are equal (within tolerance)")
74+
# return all_equal
75+
76+
import jax
77+
import jax.numpy as jnp
78+
import logging
79+
80+
def maybe_unreplicate(pytree):
81+
"""If leading axis matches device count, strip it assuming it's pmap replication."""
82+
num_devices = jax.device_count()
83+
return jax.tree_util.tree_map(
84+
lambda x: x[0] if isinstance(x, jax.Array) and x.shape[0] == num_devices else x,
85+
pytree
86+
)
87+
88+
def move_to_cpu(tree):
89+
return jax.tree_util.tree_map(lambda x: jax.device_put(x, device=jax.devices("cpu")[0]), tree)
90+
91+
4692
def are_weights_equal(params1, params2, atol=1e-6, rtol=1e-6):
47-
"""Compares two JAX PyTrees of weights and prints where they differ."""
93+
"""Compares two JAX PyTrees of weights and logs where they differ, safely handling PMAP replication."""
94+
# Attempt to unreplicate if needed
95+
params1 = maybe_unreplicate(params1)
96+
params2 = maybe_unreplicate(params2)
97+
98+
params1 = move_to_cpu(params1)
99+
params2 = move_to_cpu(params2)
100+
48101
all_equal = True
49102

50103
def compare_fn(p1, p2):
51104
nonlocal all_equal
52-
#if not jnp.allclose(p1, p2):
53105
if not jnp.allclose(p1, p2, atol=atol, rtol=rtol):
54-
print("❌ Mismatch found:")
55-
print(f"Shape 1: {p1.shape}, Shape 2: {p2.shape}")
56-
print(f"Max diff: {jnp.max(jnp.abs(p1 - p2))}")
106+
logging.info("❌ Mismatch found:")
107+
logging.info(f"Shape 1: {p1.shape}, Shape 2: {p2.shape}")
108+
logging.info(f"Max diff: {jnp.max(jnp.abs(p1 - p2))}")
57109
all_equal = False
58110
return jnp.allclose(p1, p2, atol=atol, rtol=rtol)
59111

60112
try:
61-
_ = jax.tree_util.tree_map(compare_fn, params1, params2)
113+
jax.tree_util.tree_map(compare_fn, params1, params2)
62114
except Exception as e:
63-
print("❌ Structure mismatch or error during comparison:", e)
115+
logging.info("❌ Structure mismatch or error during comparison:", exc_info=True)
64116
return False
65117

66118
if all_equal:
67-
print("✅ All weights are equal (within tolerance)")
119+
logging.info("✅ All weights are equal (within tolerance)")
68120
return all_equal
121+
122+
123+
124+
def use_pytorch_weights2(jax_params, file_name=None, replicate=False):
125+
126+
def deep_copy_to_cpu(pytree):
127+
return tree_map(lambda x: jax.device_put(jnp.array(copy.deepcopy(x)), device=jax.devices("cpu")[0]), pytree)
128+
129+
breakpoint()
130+
jax_copy = deep_copy_to_cpu(jax_params)
131+
# Load PyTorch state_dict lazily to CPU
132+
state_dict = torch.load(file_name, map_location='cpu')
133+
print(state_dict.keys())
134+
# Convert PyTorch tensors to NumPy arrays
135+
numpy_weights = {k: v.cpu().numpy() for k, v in state_dict.items()}
136+
137+
# --- Embedding Table ---
138+
embedding_table = np.concatenate([
139+
numpy_weights[f'embedding_chunk_{i}'] for i in range(4)
140+
], axis=0) # adjust axis depending on chunking direction
141+
142+
jax_copy['embedding_table'] = jnp.array(embedding_table)
143+
144+
# --- Bot MLP: Dense_0 to Dense_2 ---
145+
for i, j in zip([0, 2, 4], range(3)):
146+
jax_copy[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'bot_mlp.{i}.weight'].T)
147+
jax_copy[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'bot_mlp.{i}.bias'])
148+
149+
# --- Top MLP: Dense_3 to Dense_7 ---
150+
for i, j in zip([0, 2, 4, 6, 8], range(3, 8)):
151+
jax_copy[f'Dense_{j}']['kernel'] = jnp.array(numpy_weights[f'top_mlp.{i}.weight'].T)
152+
jax_copy[f'Dense_{j}']['bias'] = jnp.array(numpy_weights[f'top_mlp.{i}.bias'])
153+
#jax_copy = tree_map(lambda x: jnp.array(x), jax_copy)
154+
del state_dict
155+
156+
return jax_copy
157+

‎reference_algorithms/schedule_free/jax/submission.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
import logging
1212
from optax.contrib import schedule_free_adamw
1313
from algoperf import spec
14-
from custom_pytorch_jax_converter import use_pytorch_weights, are_weights_equal
14+
from custom_pytorch_jax_converter import use_pytorch_weights_inplace, are_weights_equal, use_pytorch_weights2
1515

1616
_GRAD_CLIP_EPS = 1e-6
1717

@@ -142,6 +142,8 @@ def update_params(workload: spec.Workload,
142142
del loss_type
143143
del eval_results
144144

145+
breakpoint()
146+
145147
optimizer_state, opt_update_fn = optimizer_state
146148
per_device_rngs = jax.random.split(rng, jax.local_device_count())
147149
if hasattr(hyperparameters, 'label_smoothing'):
@@ -170,9 +172,17 @@ def update_params(workload: spec.Workload,
170172
'loss': loss[0],
171173
'grad_norm': grad_norm[0],
172174
}, global_step)
173-
logging.info('\n\nUsing the PyTorch weights of first update\n\n')
174-
params = use_pytorch_weights(new_params, file_name="/results/pytorch_base_model_criteo1tb_8_june_after_first_update.pth")
175-
are_weights_equal(new_params, params)
175+
logging.info(f'\n\nUsing the PyTorch weights of Global Step {global_step} to validate\n\n')
176+
if global_step% 100 == 0:
177+
current_date = "2025-06-14"
178+
file_name = f"/results/schedule_free_test_pytorch_weights/criteo1tb_{current_date}_after_{global_step}_steps.pth"
179+
# breakpoint()
180+
breakpoint()
181+
params = use_pytorch_weights2(new_params, file_name=file_name, replicate=True)
182+
are_weights_equal(new_params, params)
183+
del params
184+
185+
176186
return (new_optimizer_state, opt_update_fn), new_params, new_model_state
177187

178188

‎reference_algorithms/schedule_free/pytorch/submission.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,13 @@ def closure():
280280
global_step,
281281
loss.item())
282282

283-
torch.save(current_param_container.module.state_dict(), "/results/pytorch_base_model_criteo1tb_8_june_after_first_update.pth")
283+
if global_step % 100 == 0:
284+
from datetime import date, timedelta
285+
current_date = date.today().strftime("%Y-%m-%d")
286+
file_name = f"/results/schedule_free_test_pytorch_weights/criteo1tb_{current_date}_after_{global_step}_steps.pth"
287+
logging.info(f"Saving model to {file_name}")
288+
# Save the model state dict
289+
torch.save(current_param_container.module.state_dict(), file_name)
284290

285291
return (optimizer_state, current_param_container, new_model_state)
286292

@@ -331,4 +337,4 @@ def data_selection(workload: spec.Workload,
331337
del global_step
332338
del rng
333339
batch = next(input_queue)
334-
return batch
340+
return batch

0 commit comments

Comments
 (0)