Skip to content

cp: chore: Merge fp8 args (1279) into r0.2.0 #1306

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鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ko3n1g merged 1 commit into from
Nov 12, 2025
Merged

cp: chore: Merge fp8 args (1279) into r0.2.0 #1306

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
15 changes: 1 addition & 14 deletions scripts/performance/argument_parser.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,24 +134,11 @@ def parse_cli_args():
"-c",
"--compute_dtype",
type=str,
choices=["bf16", "fp8"],
choices=["bf16", "fp8_cs", "fp8_mx", "fp8_sc"],
help="Compute precision. Options- bf16 or fp8. Defaults to bf16",
required=False,
default="bf16",
)
fp8_recipe_msg = (
"FP8 recipe. Options- ds (per-tensor delayed scaling), cs (per-tensor current scaling), "
"mxfp8, ss (subchannel scaling). Defaults to ds"
)
parser.add_argument(
"-fr",
"--fp8_recipe",
type=str,
choices=["cs", "mx", "sc"],
help=fp8_recipe_msg,
required=False,
default="cs",
)
parser.add_argument(
"--task",
choices=["pretrain", "sft", "lora"],
Expand Down
24 changes: 12 additions & 12 deletions scripts/performance/configs/deepseek/deepseek_llm_pretrain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,16 +44,16 @@ def set_deepseek_v3_common_configs(cfg: ConfigContainer) -> None:
cfg.model.moe_router_force_load_balancing = True


def deepseek_v3_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def deepseek_v3_gb300_config(precision: str = "bf16") -> ConfigContainer:
"""GB300, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.DEEPSEEK_V3_GB300_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.DEEPSEEK_V3_GB300_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.DEEPSEEK_V3_GB300_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = pretrain_config(
mock=True,
Expand All @@ -76,16 +76,16 @@ def deepseek_v3_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") ->
return cfg


def deepseek_v3_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def deepseek_v3_gb200_config(precision: str = "bf16") -> ConfigContainer:
"""GB200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.DEEPSEEK_V3_GB200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.DEEPSEEK_V3_GB200_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.DEEPSEEK_V3_GB200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = pretrain_config(
mock=True,
Expand All @@ -108,16 +108,16 @@ def deepseek_v3_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") ->
return cfg


def deepseek_v3_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def deepseek_v3_b200_config(precision: str = "bf16") -> ConfigContainer:
"""B200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.DEEPSEEK_V3_B200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.DEEPSEEK_V3_B200_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.DEEPSEEK_V3_B200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = pretrain_config(
mock=True,
Expand All @@ -135,16 +135,16 @@ def deepseek_v3_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") ->
return cfg


def deepseek_v3_h100_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def deepseek_v3_h100_config(precision: str = "bf16") -> ConfigContainer:
"""H100, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.DEEPSEEK_V3_H100_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.DEEPSEEK_V3_H100_FP8_CS_BASE_CONFIG
if fp8_recipe == "sc":
if precision == "fp8_sc":
base_cfg = base_cfgs.DEEPSEEK_V3_H100_FP8_SC_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = pretrain_config(
mock=True,
Expand Down
16 changes: 8 additions & 8 deletions scripts/performance/configs/gpt_oss/gpt_oss_llm_pretrain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,14 +36,14 @@ def set_gpt_oss_common_configs(cfg: ConfigContainer) -> None:
cfg.model.moe_router_force_load_balancing = True


def gpt_oss_120b_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def gpt_oss_120b_gb300_config(precision: str = "bf16") -> ConfigContainer:
"""GB300, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.GPT_OSS_120B_GB300_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.GPT_OSS_120B_GB300_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = gpt_oss_120b_pretrain_config(
mock=True,
Expand All @@ -55,14 +55,14 @@ def gpt_oss_120b_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") -
return cfg


def gpt_oss_120b_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def gpt_oss_120b_gb200_config(precision: str = "bf16") -> ConfigContainer:
"""GB200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.GPT_OSS_120B_GB200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.GPT_OSS_120B_GB200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = gpt_oss_120b_pretrain_config(
mock=True,
Expand All @@ -74,14 +74,14 @@ def gpt_oss_120b_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") -
return cfg


def gpt_oss_120b_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def gpt_oss_120b_b200_config(precision: str = "bf16") -> ConfigContainer:
"""B200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.GPT_OSS_120B_B200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.GPT_OSS_120B_B200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = gpt_oss_120b_pretrain_config(
mock=True,
Expand All @@ -93,14 +93,14 @@ def gpt_oss_120b_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") ->
return cfg


def gpt_oss_120b_h100_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def gpt_oss_120b_h100_config(precision: str = "bf16") -> ConfigContainer:
"""H100, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.GPT_OSS_120B_H100_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.GPT_OSS_120B_H100_FP8_CS_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = gpt_oss_120b_pretrain_config(
mock=True,
Expand Down
44 changes: 22 additions & 22 deletions scripts/performance/configs/llama3/llama3_llm_pretrain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,17 +50,17 @@ def set_llama3_common_configs(cfg: ConfigContainer) -> None:
# Llama3 70B configs ---------------------------------------------------------


def llama3_70b_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_70b_gb300_config(precision: str = "bf16") -> ConfigContainer:
"""GB300, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_70B_GB300_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_bf16_b200_h8192_tp2_mbs1_seqlen8192
else:
base_cfg = base_cfgs.LLAMA3_70B_GB300_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.LLAMA3_70B_GB300_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192

cfg = llama3_70b_pretrain_config(mock=True, precision_config=precision_config)
Expand All @@ -77,17 +77,17 @@ def llama3_70b_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") ->
return cfg


def llama3_70b_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_70b_gb200_config(precision: str = "bf16") -> ConfigContainer:
"""GB200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_70B_GB200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_bf16_b200_h8192_tp2_mbs1_seqlen8192
else:
base_cfg = base_cfgs.LLAMA3_70B_GB200_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.LLAMA3_70B_GB200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192

cfg = llama3_70b_pretrain_config(mock=True, precision_config=precision_config)
Expand All @@ -104,17 +104,17 @@ def llama3_70b_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") ->
return cfg


def llama3_70b_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_70b_b200_config(precision: str = "bf16") -> ConfigContainer:
"""B200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_70B_B200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_bf16_b200_h8192_tp2_mbs1_seqlen8192
else:
base_cfg = base_cfgs.LLAMA3_70B_B200_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.LLAMA3_70B_B200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_fp8_b200_h8192_tp2_mbs1_seqlen8192

cfg = llama3_70b_pretrain_config(mock=True, precision_config=precision_config)
Expand All @@ -131,15 +131,15 @@ def llama3_70b_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> C
return cfg


def llama3_70b_h100_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_70b_h100_config(precision: str = "bf16") -> ConfigContainer:
"""H100, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_70B_H100_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192
else:
base_cfg = base_cfgs.LLAMA3_70B_H100_FP8_CS_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)
comm_overlap_cfg = userbuffers_fp8_h100_h8192_tp4_mbs1_seqlen8192

cfg = llama3_70b_pretrain_config(mock=True, precision_config=precision_config)
Expand All @@ -154,16 +154,16 @@ def llama3_70b_h100_config(precision: str = "bf16", fp8_recipe: str = "cs") -> C
# Llama3 8B configs ---------------------------------------------------------


def llama3_8b_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_8b_gb300_config(precision: str = "bf16") -> ConfigContainer:
"""GB300, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_8B_GB300_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.LLAMA3_8B_GB300_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.LLAMA3_8B_GB300_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = llama3_8b_pretrain_config(mock=True, precision_config=precision_config)
set_llama3_common_configs(cfg)
Expand All @@ -174,16 +174,16 @@ def llama3_8b_gb300_config(precision: str = "bf16", fp8_recipe: str = "cs") -> C
return cfg


def llama3_8b_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_8b_gb200_config(precision: str = "bf16") -> ConfigContainer:
"""GB200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_8B_GB200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.LLAMA3_8B_GB200_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.LLAMA3_8B_GB200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = llama3_8b_pretrain_config(mock=True, precision_config=precision_config)
set_llama3_common_configs(cfg)
Expand All @@ -194,16 +194,16 @@ def llama3_8b_gb200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> C
return cfg


def llama3_8b_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_8b_b200_config(precision: str = "bf16") -> ConfigContainer:
"""B200, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_8B_B200_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.LLAMA3_8B_B200_FP8_CS_BASE_CONFIG
if fp8_recipe == "mx":
if precision == "fp8_mx":
base_cfg = base_cfgs.LLAMA3_8B_B200_FP8_MX_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = llama3_8b_pretrain_config(mock=True, precision_config=precision_config)
set_llama3_common_configs(cfg)
Expand All @@ -214,14 +214,14 @@ def llama3_8b_b200_config(precision: str = "bf16", fp8_recipe: str = "cs") -> Co
return cfg


def llama3_8b_h100_config(precision: str = "bf16", fp8_recipe: str = "cs") -> ConfigContainer:
def llama3_8b_h100_config(precision: str = "bf16") -> ConfigContainer:
"""H100, baseline config."""
if precision == "bf16":
base_cfg = base_cfgs.LLAMA3_8B_H100_BF16_BASE_CONFIG
precision_config = get_precision_config(precision)
else:
base_cfg = base_cfgs.LLAMA3_8B_H100_FP8_CS_BASE_CONFIG
precision_config = get_precision_config(precision, fp8_recipe)
precision_config = get_precision_config(precision)

cfg = llama3_8b_pretrain_config(mock=True, precision_config=precision_config)
set_llama3_common_configs(cfg)
Expand Down
Loading