SimuMax: LLM Distributed Training Simulator

To ensure a legal configuration, the parallel dimensions must satisfy specific divisibility rules based on the model type.

For Dense Models: The data parallel size (dp) is calculated as: dp = world_size / (tp * pp * cp)

Constraint: world_size must be divisible by (tp * pp * cp).

For MoE Models: SimuMax also checks: world_size % (ep * etp * pp) == 0

Practical Rule:

1. Establish a legal dense tp/pp/cp split first.
2. Add ep (Expert Parallelism).
3. Verify model-specific expert divisibility (e.g., expert_num % ep == 0).

You can use shipped configurations directly if your target machine and dominant operator shapes are already covered.

However, you should measure and add your own system data if:

• The target machine is new to the simulator.
• The system.miss_efficiency field is non-empty for your target case.

Accurate machine measurement is required before interpreting timing results for new hardware.

For formal B200 real memory results, the following measurement contract is used:

• real alloc = maximum over all ranks of torch.cuda.max_memory_allocated()
• real reserved = maximum over all ranks of torch.cuda.max_memory_reserved()

Important:

• Partial-rank memory logging is considered invalid for formal comparison.
• For memory refresh runs, use 2 steps and 0 warmup steps.
• If only mbc changes for the same topology, one representative real-memory run may be reused.

The strategy file is a core input for SimuMax that defines the training runtime configuration. It aligns closely with Megatron runtime semantics. To ensure accurate timing and memory estimations, the settings for TP/PP/EP, sequence parallel, recompute, and VPP in the strategy file must match your actual real run.

SimuMax requires three core input files to function:

1. system (hardware configuration)
2. strategy (runtime/parallelism configuration)
3. model (model architecture configuration)

SimuMax uses a model configuration file (located in configs/models) to describe the static structure of an LLM. This file is the foundation for modeling cost, memory, and simulation. A complete SimuMax run requires three input files working together:

1. system: Machine capability and efficiency data.
2. strategy: Parallelism and runtime strategies.
3. model: The model structure description.

To ensure accurate timing and memory estimation, the model configuration must align with the target real-world model (e.g., Megatron-based models).

SimuMax uses a system.json file to describe the hardware environment. A complete configuration consists of three logical parts:

1. Basic Information: Includes the system name (sys_name) and the number of GPUs per node (num_per_node).
2. accelerator: Defines compute capability (TFLOPS), memory capacity (mem_gbs), and memory bandwidth behavior.
3. networks: Defines intra-node and inter-node communication behavior, including bandwidth, latency, and operator-specific scaling factors.

Some machine families may include additional specific fields (e.g., FC8), but the three parts above form the core structure.

SimuMax can generate simulator traces and memory artifacts to help you understand:

• Stage-level timing
• Memory peaks
• Cache lifetime

These artifacts are useful for comparing modeled results against real traces or real memory evidence.

Note: If you only need a smoke test or a rough OOM (Out of Memory) estimate, you do not need to enable trace or snapshot generation.

SimuMax provides shipped configurations for common models and systems.

Use shipped configs when:

• The target machine is similar to a provided example.
• The model uses standard operator shapes already covered.
• You need quick OOM feasibility or rough exploration.

Measure your own machine/model when:

• The hardware is new or communication bandwidth/latency is unknown.
• The target model uses operator shapes not covered in existing configs.
• system.miss_efficiency is non-empty and you need to explain timing discrepancies.
• You require benchmark-grade timing accuracy.

SimuMax requires three core input files: system, strategy, and model. The system file describes the hardware capabilities of the machine, including accelerator compute power, memory bandwidth, intra/inter-node communication bandwidth and latency, and operator efficiency at specific shapes.

A complete system.json file consists of three main parts:

1. Basic Information: Includes sys_name and num_per_node (number of cards per node).
2. accelerator: Describes computation and memory access capabilities.
3. networks: Describes intra-node and inter-node communication.

Note: Some machine families may include additional fields like FC8, but the three parts above are the core requirements.

While shipped configurations are sufficient for demos and OOM feasibility screening, you should use the machine-measurement workflow (simu_tools/efficency_test/) when timing accuracy is critical.

Measure your own data if:

• The machine is new or significantly different from shipped configs.
• Interconnect bandwidth or latency is unknown.
• The model uses dominant matmul/attention shapes not covered by current configs.
• You need to explain timing discrepancies (i.e., when system.miss_efficiency is non-empty).

SimuMax uses JSON files located in configs/models/ to define the static architecture of a model. This architecture serves as the foundation for cost, memory, and simulator analysis.

To perform a complete analysis, SimuMax requires three distinct input files working together:

1. system: Contains machine capability and efficiency data.
2. strategy: Defines parallelism and runtime policies.
3. model: The architecture description (the JSON file).

Ensuring the model config matches the structural shape of your target (e.g., a Megatron run) is critical. If fields like layer count or attention type are mismatched, timing and memory analysis will drift even if your system and strategy configs are correct.