diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_ckpt.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_ckpt.py index 7b6774b992..20d5f0c473 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_ckpt.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_ckpt.py @@ -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. @@ -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. @@ -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": { @@ -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") @@ -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) @@ -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) @@ -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 @@ -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) @@ -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) @@ -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) @@ -312,7 +322,7 @@ 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 @@ -320,12 +330,13 @@ def convert_to_jax_weights(base_model_path: str, model_size: str): 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__": @@ -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() @@ -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 diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_unscanned_ckpt.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_unscanned_ckpt.py index f3cdda78b5..a25d7df6ae 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_unscanned_ckpt.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt_oss_unscanned_ckpt.py @@ -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. @@ -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. @@ -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": { @@ -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") @@ -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) @@ -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 ) @@ -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) @@ -299,7 +310,7 @@ 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 @@ -307,12 +318,13 @@ def convert_to_jax_weights(base_model_path: str, model_size: str): 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__": @@ -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: @@ -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,