129 lines
4.4 KiB
Python
129 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
def force_megatron_cpu_mode() -> None:
|
|
# Megatron-Core allocates several modules with device=torch.cuda.current_device().
|
|
# Returning "cpu" lets those constructors allocate CPU tensors for this probe.
|
|
torch.cuda.current_device = lambda: "cpu" # type: ignore[assignment]
|
|
torch.cuda.set_device = lambda *_args, **_kwargs: None # type: ignore[assignment]
|
|
|
|
|
|
def init_single_process() -> None:
|
|
force_megatron_cpu_mode()
|
|
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
|
os.environ.setdefault("MASTER_PORT", "29621")
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
if not dist.is_initialized():
|
|
dist.init_process_group("gloo", rank=0, world_size=1)
|
|
|
|
from megatron.core import parallel_state
|
|
|
|
if not parallel_state.is_initialized():
|
|
parallel_state.initialize_model_parallel(
|
|
tensor_model_parallel_size=1,
|
|
pipeline_model_parallel_size=1,
|
|
virtual_pipeline_model_parallel_size=None,
|
|
context_parallel_size=1,
|
|
expert_model_parallel_size=1,
|
|
)
|
|
|
|
|
|
def build_model(seq_length: int):
|
|
import torch.nn.functional as F
|
|
from megatron.bridge.models.gpt_provider import GPTModelProvider
|
|
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
|
|
from megatron.core.transformer.torch_norm import WrappedTorchNorm
|
|
from megatron.core.utils import init_method_normal, scaled_init_method_normal
|
|
|
|
def cpu_rmsnorm_block_spec(config, vp_stage=None):
|
|
spec = get_gpt_decoder_block_spec(
|
|
config,
|
|
use_transformer_engine=False,
|
|
normalization="RMSNorm",
|
|
vp_stage=vp_stage,
|
|
)
|
|
spec.layer_norm = WrappedTorchNorm
|
|
return spec
|
|
|
|
from megatron.core.process_groups_config import ProcessGroupCollection
|
|
|
|
provider = GPTModelProvider(
|
|
num_layers=12,
|
|
hidden_size=1536,
|
|
num_attention_heads=24,
|
|
num_query_groups=4,
|
|
kv_channels=64,
|
|
ffn_hidden_size=4608,
|
|
seq_length=seq_length,
|
|
vocab_size=154856,
|
|
should_pad_vocab=True,
|
|
share_embeddings_and_output_weights=True,
|
|
position_embedding_type="rope",
|
|
normalization="RMSNorm",
|
|
gated_linear_unit=True,
|
|
activation_func=F.silu,
|
|
add_bias_linear=False,
|
|
num_moe_experts=12,
|
|
moe_layer_freq=[1 if idx in {2, 4, 6, 8, 10} else 0 for idx in range(12)],
|
|
moe_ffn_hidden_size=6144,
|
|
moe_router_topk=4,
|
|
moe_router_load_balancing_type="aux_loss",
|
|
moe_aux_loss_coeff=0.02,
|
|
moe_z_loss_coeff=0.001,
|
|
moe_token_dispatcher_type="allgather",
|
|
moe_expert_capacity_factor=1.25,
|
|
moe_router_enable_expert_bias=False,
|
|
moe_router_bias_update_rate=0.02,
|
|
moe_grouped_gemm=True,
|
|
tensor_model_parallel_size=1,
|
|
pipeline_model_parallel_size=1,
|
|
expert_model_parallel_size=1,
|
|
context_parallel_size=1,
|
|
sequence_parallel=False,
|
|
transformer_impl="local",
|
|
attention_backend="unfused",
|
|
transformer_layer_spec=cpu_rmsnorm_block_spec,
|
|
init_method=init_method_normal(0.02),
|
|
output_layer_init_method=scaled_init_method_normal(0.02, 12),
|
|
init_method_std=0.02,
|
|
)
|
|
provider.perform_initialization = False
|
|
provider._pg_collection = ProcessGroupCollection.use_mpu_process_groups()
|
|
return provider.provide(pre_process=True, post_process=True)
|
|
|
|
|
|
def inspect_model(args: argparse.Namespace) -> None:
|
|
init_single_process()
|
|
model = build_model(args.seq_length)
|
|
print("model_built", type(model))
|
|
print("param_count", sum(p.numel() for p in model.parameters()))
|
|
if args.mode == "sharded-keys":
|
|
state = model.sharded_state_dict(prefix="")
|
|
print("sharded_keys_sample")
|
|
else:
|
|
state = model.state_dict()
|
|
print("state_keys_sample")
|
|
for key in list(state.keys())[:120]:
|
|
print(key)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--seq-length", type=int, default=8192)
|
|
parser.add_argument("--mode", choices=["state-keys", "sharded-keys"], default="state-keys")
|
|
args = parser.parse_args()
|
|
inspect_model(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|