Skip to content

feat: configure standalone checkpoint converters with --pure-nnx flag #4560

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
RexBearIU wants to merge 1 commit into from
+45 −20
Closed

feat: configure standalone checkpoint converters with --pure-nnx flag #4560

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 24 additions & 12 deletions src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_ckpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@


def _convert_huggingface_to_jax_weights(
base_model_path: str, model_size: str, model_params: dict, mem_info: psutil.Process
base_model_path: str, model_size: str, model_params: dict, mem_info: psutil.Process, pure_nnx: bool = True
):
"""Convert a Huggingface Checkpoint to a dictionary of Numpy arrays representing the weights.

Expand All @@ -62,6 +62,7 @@ def _convert_huggingface_to_jax_weights(
model_size (str): Size of the base model.
model_params (dict): Dictionary containing model parameters.
mem_info (psutil.Process): Process object to track memory usage.
pure_nnx (bool): Whether to format weights for pure NNX.

Returns:
jax_weights (dict): Dictionary containing the converted weights.
Expand Down Expand Up @@ -96,12 +97,14 @@ def _convert_huggingface_to_jax_weights(
"decoder": {
"decoder_norm": {"scale": None},
"logits_dense": {"kernel": None},
"layers": {},
},
}
if pure_nnx:
jax_weights["decoder"]["layers"] = {}

# block 0, 1
for block_idx in range(layer_cycle_interval):
jax_weights["decoder"]["layers"][f"layers_{block_idx}"] = {
layer_dict = {
"pre_self_attention_layer_norm": {"scale": None},
"post_self_attention_layer_norm": {"scale": None},
"GptOssAttention": {
Expand All @@ -121,6 +124,13 @@ def _convert_huggingface_to_jax_weights(
"wo_bias": None,
},
}
if pure_nnx:
jax_weights["decoder"]["layers"][block_idx] = layer_dict
else:
jax_weights["decoder"][f"layers_{block_idx}"] = layer_dict

def get_layer(b_idx):
return jax_weights["decoder"]["layers"][b_idx] if pure_nnx else jax_weights["decoder"][f"layers_{b_idx}"]

# decoder norm scale ###########################################
max_logging.log("Processing decoder norm scale")
Expand All @@ -147,7 +157,7 @@ def _convert_huggingface_to_jax_weights(
for layer_idx in MemoryMonitorTqdm(range(base_num_decoder_layers), desc="layers", leave=True):
block_layer_idx, block_idx = divmod(layer_idx, layer_cycle_interval)
stack_shape = (base_num_decoder_layers // layer_cycle_interval,)
self_attention = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]["GptOssAttention"]
self_attention = get_layer(block_idx)["GptOssAttention"]

wq = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.attention.wq.weight"], cast_dtype=CAST_DTYPE)
wk = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.attention.wk.weight"], cast_dtype=CAST_DTYPE)
Expand Down Expand Up @@ -200,7 +210,7 @@ def _convert_huggingface_to_jax_weights(
self_attention["sinks"][block_layer_idx, ...] = sinks

for block_idx in range(layer_cycle_interval):
self_attention = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]["GptOssAttention"]
self_attention = get_layer(block_idx)["GptOssAttention"]
self_attention["query"]["kernel"] = self_attention["query"]["kernel"].transpose(1, 0, 2, 3)
self_attention["key"]["kernel"] = self_attention["key"]["kernel"].transpose(1, 0, 2, 3)
self_attention["value"]["kernel"] = self_attention["value"]["kernel"].transpose(1, 0, 2, 3)
Expand All @@ -218,7 +228,7 @@ def _convert_huggingface_to_jax_weights(
for layer_idx in MemoryMonitorTqdm(range(base_num_decoder_layers), desc="layers", leave=True):
block_layer_idx, block_idx = divmod(layer_idx, layer_cycle_interval)
stack_shape = (base_num_decoder_layers // layer_cycle_interval,)
layer_weight = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]
layer_weight = get_layer(block_idx)

pre_self_attention_layernorm = _pt_to_np(
chkpt_vars[f"layers.{layer_idx}.attention_norm.weight"], cast_dtype=CAST_DTYPE
Expand All @@ -237,7 +247,7 @@ def _convert_huggingface_to_jax_weights(
layer_weight["post_self_attention_layer_norm"]["scale"][block_layer_idx, ...] = post_self_attention_layernorm # pylint: disable=E1137

for block_idx in range(layer_cycle_interval):
layer_weight = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]
layer_weight = get_layer(block_idx)
layer_weight["pre_self_attention_layer_norm"]["scale"] = layer_weight["pre_self_attention_layer_norm"][
"scale"
].transpose(1, 0)
Expand All @@ -253,7 +263,7 @@ def _convert_huggingface_to_jax_weights(
for layer_idx in MemoryMonitorTqdm(range(base_num_decoder_layers), desc="layers", leave=True):
block_layer_idx, block_idx = divmod(layer_idx, layer_cycle_interval)
stack_shape = (base_num_decoder_layers // layer_cycle_interval,)
mlp_weight = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]["GptOssMlp"]
mlp_weight = get_layer(block_idx)["GptOssMlp"]

gate = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.feed_forward.gate.weight"], cast_dtype=CAST_DTYPE)
gate_bias = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.feed_forward.gate.bias"], cast_dtype=CAST_DTYPE)
Expand Down Expand Up @@ -296,7 +306,7 @@ def _convert_huggingface_to_jax_weights(
gc.collect()

for block_idx in range(layer_cycle_interval):
mlp_weight = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]["GptOssMlp"]
mlp_weight = get_layer(block_idx)["GptOssMlp"]
mlp_weight["gate"]["kernel"] = mlp_weight["gate"]["kernel"].transpose(1, 0, 2)
mlp_weight["wi_0"] = mlp_weight["wi_0"].transpose(1, 0, 2, 3)
mlp_weight["wi_1"] = mlp_weight["wi_1"].transpose(1, 0, 2, 3)
Expand All @@ -312,20 +322,21 @@ def _convert_huggingface_to_jax_weights(
return jax_weights


def convert_to_jax_weights(base_model_path: str, model_size: str):
def convert_to_jax_weights(base_model_path: str, model_size: str, pure_nnx: bool = True):
"""
Function to convert the checkpoint at base_model_path into Orbax checkpoint
for MaxText and output jax_weights ready for MaxText

Args:
base_model_path: checkpoint path
model_size: gpt-oss-20b, gpt-oss-120b
pure_nnx: whether to format weights for pure NNX
"""
model_params = MODEL_PARAMS_DICT[model_size]
mem_info = psutil.Process()
logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
max_logging.log(f"Loading the base model from {base_model_path}")
return _convert_huggingface_to_jax_weights(base_model_path, model_size, model_params, mem_info)
return _convert_huggingface_to_jax_weights(base_model_path, model_size, model_params, mem_info, pure_nnx)


if __name__ == "__main__":
Expand All @@ -336,6 +347,7 @@ def convert_to_jax_weights(base_model_path: str, model_size: str):
parser.add_argument("--simulated-cpu-devices-count", type=int, required=False, default=16)
parser.add_argument("--use-ocdbt", type=str2bool, required=False, default=True)
parser.add_argument("--use-zarr3", type=str2bool, required=False, default=True)
parser.add_argument("--pure-nnx", type=str2bool, required=False, default=True)
args = parser.parse_args()

overall_start = time.time()
Expand All @@ -348,7 +360,7 @@ def convert_to_jax_weights(base_model_path: str, model_size: str):

# transform
start = time.time()
weights = convert_to_jax_weights(args.base_model_path, args.model_size)
weights = convert_to_jax_weights(args.base_model_path, args.model_size, args.pure_nnx)
max_logging.log(f"Elapse for transform: {(time.time() - start) / 60:.2f} min")

# save
Expand Down
29 changes: 21 additions & 8 deletions src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_unscanned_ckpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ def _hf_to_maxtext_mapping(layer_idx: int = -1) -> dict:


def _convert_huggingface_to_jax_weights(
base_model_path: str, model_size: str, model_params: dict, mem_info: psutil.Process
base_model_path: str, model_size: str, model_params: dict, mem_info: psutil.Process, pure_nnx: bool = True
):
"""Convert a Huggingface Checkpoint to a dictionary of Numpy arrays representing the weights.

Expand All @@ -129,6 +129,7 @@ def _convert_huggingface_to_jax_weights(
model_size (str): Size of the base model.
model_params (dict): Dictionary containing model parameters.
mem_info (psutil.Process): Process object to track memory usage.
pure_nnx (bool): Whether to format weights for pure NNX.

Returns:
jax_weights (dict): Dictionary containing the converted weights.
Expand Down Expand Up @@ -164,8 +165,11 @@ def _convert_huggingface_to_jax_weights(
"logits_dense": {"kernel": None},
},
}
if pure_nnx:
jax_weights["decoder"]["layers"] = {}

for layer_idx in range(base_num_decoder_layers):
jax_weights["decoder"][f"layers_{layer_idx}"] = {
layer_dict = {
"pre_self_attention_layer_norm": {"scale": None},
"post_self_attention_layer_norm": {"scale": None},
"GptOssAttention": {
Expand All @@ -185,6 +189,13 @@ def _convert_huggingface_to_jax_weights(
"wo_bias": None,
},
}
if pure_nnx:
jax_weights["decoder"]["layers"][layer_idx] = layer_dict
else:
jax_weights["decoder"][f"layers_{layer_idx}"] = layer_dict

def get_layer(l_idx):
return jax_weights["decoder"]["layers"][l_idx] if pure_nnx else jax_weights["decoder"][f"layers_{l_idx}"]

# decoder norm scale ###########################################
max_logging.log("Processing decoder norm scale")
Expand All @@ -209,7 +220,7 @@ def _convert_huggingface_to_jax_weights(
# self attention ###############################################
max_logging.log("Processing self attention")
for layer_idx in tqdm(range(base_num_decoder_layers), desc="layers", leave=False):
self_attention = jax_weights["decoder"][f"layers_{layer_idx}"]["GptOssAttention"]
self_attention = get_layer(layer_idx)["GptOssAttention"]

wq = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.attention.wq.weight"], cast_dtype=CAST_DTYPE)
wk = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.attention.wk.weight"], cast_dtype=CAST_DTYPE)
Expand Down Expand Up @@ -247,7 +258,7 @@ def _convert_huggingface_to_jax_weights(
# layer weight pre and post self attention norm ################
max_logging.log("Processing pre and post self attention norms")
for layer_idx in tqdm(range(base_num_decoder_layers), desc="layers", leave=False):
layer_weight = jax_weights["decoder"][f"layers_{layer_idx}"]
layer_weight = get_layer(layer_idx)
pre_self_attention_layernorm = _pt_to_np(
chkpt_vars[f"layers.{layer_idx}.attention_norm.weight"], cast_dtype=CAST_DTYPE
)
Expand All @@ -262,7 +273,7 @@ def _convert_huggingface_to_jax_weights(
max_logging.log("Processing layer weights")

for layer_idx in tqdm(range(base_num_decoder_layers), desc="layers", leave=False):
mlp_weight = jax_weights["decoder"][f"layers_{layer_idx}"]["GptOssMlp"]
mlp_weight = get_layer(layer_idx)["GptOssMlp"]

gate = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.feed_forward.gate.weight"], cast_dtype=CAST_DTYPE)
gate_bias = _pt_to_np(chkpt_vars[f"layers.{layer_idx}.feed_forward.gate.bias"], cast_dtype=CAST_DTYPE)
Expand Down Expand Up @@ -299,20 +310,21 @@ def _convert_huggingface_to_jax_weights(
return jax_weights


def convert_to_jax_weights(base_model_path: str, model_size: str):
def convert_to_jax_weights(base_model_path: str, model_size: str, pure_nnx: bool = True):
"""
Function to convert the checkpoint at base_model_path into Orbax checkpoint
for MaxText and output jax_weights ready for MaxText

Args:
base_model_path: checkpoint path
model_size: gpt-oss-20b, gpt-oss-120b
pure_nnx: whether to format weights for pure NNX
"""
model_params = MODEL_PARAMS_DICT[model_size]
mem_info = psutil.Process()
logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
max_logging.log(f"Loading the base model from {base_model_path}")
return _convert_huggingface_to_jax_weights(base_model_path, model_size, model_params, mem_info)
return _convert_huggingface_to_jax_weights(base_model_path, model_size, model_params, mem_info, pure_nnx)


if __name__ == "__main__":
Expand All @@ -323,6 +335,7 @@ def convert_to_jax_weights(base_model_path: str, model_size: str):
parser.add_argument("--simulated-cpu-devices-count", type=int, required=False, default=16)
parser.add_argument("--use-ocdbt", type=str2bool, required=False, default=True)
parser.add_argument("--use-zarr3", type=str2bool, required=False, default=True)
parser.add_argument("--pure-nnx", type=str2bool, required=False, default=True)
args = parser.parse_args()

if args.model_size not in MODEL_PARAMS_DICT:
Expand All @@ -333,7 +346,7 @@ def convert_to_jax_weights(base_model_path: str, model_size: str):

save_weights_to_checkpoint(
args.maxtext_model_path,
convert_to_jax_weights(args.base_model_path, args.model_size),
convert_to_jax_weights(args.base_model_path, args.model_size, args.pure_nnx),
args.simulated_cpu_devices_count,
args.use_ocdbt,
args.use_zarr3,
Expand Down
Loading