diff --git a/docs-next/public/llms.txt b/docs-next/public/llms.txt index fe67a831f..2e8b56bfd 100644 --- a/docs-next/public/llms.txt +++ b/docs-next/public/llms.txt @@ -102,6 +102,7 @@ Repo: https://github.com/modal-projects/modal-dojo - [Gemma4_26B_A4B_Recipe](https://dojo.modal.dev/reference/gemma4_26b_a4b_recipe) - [GLM_4_7_Recipe](https://dojo.modal.dev/reference/glm_4_7_recipe) - [GLM_5_3_Flash_LoRA_Recipe](https://dojo.modal.dev/reference/glm_5_3_flash_lora_recipe) +- [GLM_5_3_Flash_Recipe](https://dojo.modal.dev/reference/glm_5_3_flash_recipe) - [Inkling_Small_LoRA_Recipe](https://dojo.modal.dev/reference/inkling_small_lora_recipe) - [Inkling_Small_Recipe](https://dojo.modal.dev/reference/inkling_small_recipe) - [Kimi_K3_LoRA_Recipe](https://dojo.modal.dev/reference/kimi_k3_lora_recipe) diff --git a/modal_dojo/__init__.py b/modal_dojo/__init__.py index 12985e77d..a7a4873f7 100644 --- a/modal_dojo/__init__.py +++ b/modal_dojo/__init__.py @@ -61,6 +61,10 @@ "GLM_4_7_Recipe", ), "GLM_5_3_Flash": ("modal_dojo.common.models", "GLM_5_3_Flash"), + "GLM_5_3_Flash_Recipe": ( + "modal_dojo.train_recipes.miles_recipe", + "GLM_5_3_Flash_Recipe", + ), "GLM_5_3_Flash_LoRA": ("modal_dojo.common.models", "GLM_5_3_Flash_LoRA"), "GLM_5_3_Flash_LoRA_Recipe": ( "modal_dojo.train_recipes.miles_recipe", @@ -213,6 +217,7 @@ "GLM_4_7", "GLM_4_7_Recipe", "GLM_5_3_Flash", + "GLM_5_3_Flash_Recipe", "GLM_5_3_Flash_LoRA", "GLM_5_3_Flash_LoRA_Recipe", "HarborDataset", diff --git a/modal_dojo/common/models/validation.py b/modal_dojo/common/models/validation.py index 929e9ab18..63995ec57 100644 --- a/modal_dojo/common/models/validation.py +++ b/modal_dojo/common/models/validation.py @@ -13,7 +13,7 @@ from .deepseek_v41_flash import DeepSeek_V4_1_Flash from .gemma4_26b_a4b import Gemma4_26B_A4B from .glm_4_7 import GLM_4_7 -from .glm_5_3_flash import GLM_5_3_Flash_LoRA +from .glm_5_3_flash import GLM_5_3_Flash, GLM_5_3_Flash_LoRA from .inkling_small import Inkling_Small, Inkling_Small_LoRA from .kimi_k3 import Kimi_K3 from .moonlight_16b_a3b_instruct import Moonlight_16B_A3B_Instruct @@ -84,6 +84,7 @@ def find(cls, name: str) -> "_ValidationConfig": VALIDATION_CONFIGS: set[_ValidationConfig] = { _ValidationConfig("GLM-4.7", GLM_4_7, Framework.SLIME), + _ValidationConfig("GLM-5.3-Flash", GLM_5_3_Flash, Framework.MILES), _ValidationConfig("GLM-5.3-Flash-LoRA", GLM_5_3_Flash_LoRA, Framework.MILES), _ValidationConfig("Qwen3-0.6B", Qwen3_0_6B, Framework.SLIME), _ValidationConfig("Qwen3-1.7B", Qwen3_1_7B, Framework.SLIME), diff --git a/modal_dojo/frameworks/miles/modal_helpers/patches/patch_glm_5_3_flash_fp8_device.py b/modal_dojo/frameworks/miles/modal_helpers/patches/patch_glm_5_3_flash_fp8_device.py new file mode 100644 index 000000000..66fbdb642 --- /dev/null +++ b/modal_dojo/frameworks/miles/modal_helpers/patches/patch_glm_5_3_flash_fp8_device.py @@ -0,0 +1,36 @@ +""" +image: radixark/miles:glm53next +file: mbridge/models/ext/deepseek_v3/dequant_fp8_safetensor_io.py + +Keep safetensors reads on each torchrun rank's GPU; bare cuda selects GPU 0. +""" + +from pathlib import Path + +MARKER = "PATCHED_MODAL_DOJO_GLM53_FP8_DEVICE" + + +def patch_source(source: str) -> str: + if MARKER in source: + return source + old = 'device="cuda"' + if source.count(old) != 2: + raise ValueError("GLM FP8 reader changed: expected two CUDA device anchors") + # safetensors interprets bare 'cuda' as cuda:0, independent of set_device. + source = source.replace(old, 'device=f"cuda:{torch.cuda.current_device()}"') + return f"# {MARKER}\n" + source + + +def main() -> None: + from importlib.util import find_spec + + path = ( + Path(find_spec("mbridge").origin).parent + / "models/ext/deepseek_v3/dequant_fp8_safetensor_io.py" + ) + path.write_text(patch_source(path.read_text())) + print(f"Applied {MARKER} to {path}") + + +if __name__ == "__main__": + main() diff --git a/modal_dojo/frameworks/miles/modal_helpers/patches/patch_glm_5_3_flash_timing.py b/modal_dojo/frameworks/miles/modal_helpers/patches/patch_glm_5_3_flash_timing.py new file mode 100644 index 000000000..cb5c10653 --- /dev/null +++ b/modal_dojo/frameworks/miles/modal_helpers/patches/patch_glm_5_3_flash_timing.py @@ -0,0 +1,171 @@ +""" +image: radixark/miles:glm53next +commit: https://github.com/radixark/miles/commit/cc76e23915b2132ecfff65b97331fd83b02637e6 +file: train.py, miles/backends/megatron_utils/{actor,model}.py + +Adapt Dojo step and substep timing to this driver; the shared patch targets +an older Miles API. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +MARKER = "PATCHED_MODAL_DOJO_GLM53_TIMING" +PREAMBLE = f"""# {MARKER} +import sys as _dojo_glm_sys +if '/root' not in _dojo_glm_sys.path: + _dojo_glm_sys.path.insert(0, '/root') +from modal_dojo.common.timing_recorder import ( + recording_lane as _dojo_glm_lane, + recording_lane_on_reporting_rank as _dojo_glm_actor, + time_phase as _dojo_glm_phase, +) +from modal_dojo.frameworks.miles.phase_reporting import report_step_event as _dojo_glm_report +""" + + +def _wrap(source: str, spans: list[tuple[int, int, str]]) -> str: + lines = source.splitlines(keepends=True) + for start, end, header in sorted(spans, reverse=True): + indent = lines[start - 1][ + : len(lines[start - 1]) - len(lines[start - 1].lstrip()) + ] + body = [ + " " + line if line.strip() else line for line in lines[start - 1 : end] + ] + lines[start - 1 : end] = [ + *(indent + line + "\n" for line in header.splitlines()), + *body, + ] + return "".join(lines) + + +def _function(source: str, name: str): + matches = [ + node + for node in ast.walk(ast.parse(source)) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == name + ] + if len(matches) != 1: + raise ValueError(f"expected one {name} function, found {len(matches)}") + return matches[0] + + +def _wrap_function(source: str, name: str, header: str) -> str: + node = _function(source, name) + body = node.body + if isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): + body = body[1:] + return _wrap(source, [(body[0].lineno, body[-1].end_lineno, header)]) + + +def _wrap_calls(source: str, mapping: dict[str, str]) -> str: + spans = [] + found = set() + for node in ast.walk(ast.parse(source)): + if not isinstance(node, (ast.Expr, ast.Assign, ast.Return)): + continue + value = node.value + if isinstance(value, ast.Await): + value = value.value + if isinstance(value, ast.Call) and (name := ast.unparse(value.func)) in mapping: + found.add(name) + header = mapping[name] + if name == "update_weights" and not value.keywords: + header = "with _dojo_glm_lane('driver', None), _dojo_glm_phase('initial_weight_sync'):" + spans.append((node.lineno, node.end_lineno, header)) + if missing := mapping.keys() - found: + raise ValueError(f"missing GLM timing calls: {sorted(missing)}") + return _wrap(source, spans) + + +def patch_source(source: str, target: str) -> str: + if MARKER in source: + return source + if target == "train.py": + source = _wrap_calls( + source, + { + "create_rollout_components": "with _dojo_glm_lane('driver', None), _dojo_glm_phase('initialize_rollouts'):", + "create_training_models": "with _dojo_glm_lane('driver', None), _dojo_glm_phase('initialize_train'):", + "rollout_executor.get.remote": "with _dojo_glm_phase('generate_rollouts'):\n _dojo_glm_report('generate_rollouts', args, rollout_id)", + "actor_model.train": "with _dojo_glm_phase('train_models'):", + "inference_controller.offload": "with _dojo_glm_phase('offload_rollout'):", + "offload_train": "with _dojo_glm_phase('offload_train'):", + "save": "with _dojo_glm_phase('checkpoint_save'):", + "update_weights": "with _dojo_glm_phase('weight_sync'):", + "eval_dispatcher.dispatch": "with _dojo_glm_phase('evaluate_rollouts'):", + }, + ) + tree = ast.parse(source) + loops = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.For) and ast.unparse(node.target) == "rollout_id" + ] + if len(loops) != 1: + raise ValueError("expected one driver rollout loop") + loop = loops[0] + source = _wrap( + source, + [ + ( + loop.body[0].lineno, + loop.body[-1].end_lineno, + "with _dojo_glm_lane('driver', rollout_id):", + ) + ], + ) + elif target == "actor.py": + source = _wrap_function( + source, "compute_log_prob", "with _dojo_glm_phase('compute_log_probs'):" + ) + source = _wrap_function( + source, + "train", + "with _dojo_glm_actor(rollout_id, 'critic' if self.role == 'critic' else 'actor'):", + ) + elif target == "model.py": + source = _wrap_calls( + source, + { + "run_forward_backward_pass": "with _dojo_glm_phase('forward_backward'):", + "optimizer.step": "with _dojo_glm_phase('optimizer_step'):", + }, + ) + else: + raise ValueError(f"unknown GLM timing target: {target}") + lines = source.splitlines(keepends=True) + # Keep module docstrings and future imports at the top. + insertion = 0 + for node in ast.parse(source).body: + if (isinstance(node, ast.ImportFrom) and node.module == "__future__") or ( + isinstance(node, ast.Expr) + and isinstance(node.value, ast.Constant) + and isinstance(node.value.value, str) + ): + insertion = node.end_lineno + else: + break + lines.insert(insertion, PREAMBLE) + source = "".join(lines) + compile(source, target, "exec") + return source + + +def main(root: Path = Path("/root/miles")) -> None: + for relative in ( + "train.py", + "miles/backends/megatron_utils/actor.py", + "miles/backends/megatron_utils/model.py", + ): + path = root / relative + path.write_text(patch_source(path.read_text(), path.name)) + print(f"Applied {MARKER} to {path}") + + +if __name__ == "__main__": + main() diff --git a/modal_dojo/train_recipes/miles_recipe/__init__.py b/modal_dojo/train_recipes/miles_recipe/__init__.py index 1d7b3353f..cf309b2df 100644 --- a/modal_dojo/train_recipes/miles_recipe/__init__.py +++ b/modal_dojo/train_recipes/miles_recipe/__init__.py @@ -5,6 +5,9 @@ from modal_dojo.train_recipes.miles_recipe.gemma4_26b_a4b import ( Gemma4_26B_A4B_Recipe, ) +from modal_dojo.train_recipes.miles_recipe.glm_5_3_flash import ( + GLM_5_3_Flash_Recipe, +) from modal_dojo.train_recipes.miles_recipe.glm_5_3_flash_lora import ( GLM_5_3_Flash_LoRA_Recipe, ) @@ -26,6 +29,7 @@ "MilesRecipe", "DeepSeek_V4_1_Flash_Recipe", "Gemma4_26B_A4B_Recipe", + "GLM_5_3_Flash_Recipe", "GLM_5_3_Flash_LoRA_Recipe", "Inkling_Small_Recipe", "Inkling_Small_LoRA_Recipe", diff --git a/modal_dojo/train_recipes/miles_recipe/glm_5_3_flash.py b/modal_dojo/train_recipes/miles_recipe/glm_5_3_flash.py new file mode 100644 index 000000000..e89b23b1b --- /dev/null +++ b/modal_dojo/train_recipes/miles_recipe/glm_5_3_flash.py @@ -0,0 +1,153 @@ +"""GLM-5.3-Flash text GRPO, based on radixark/miles#2786.""" + +from dataclasses import field +from pathlib import Path +from typing import Any, ClassVar + +from pydantic import ConfigDict, model_validator +from pydantic.dataclasses import dataclass + +from modal_dojo.common.models.glm_5_3_flash import GLM_5_3_Flash +from modal_dojo.common.models.base import ModelConfig +from modal_dojo.common.patches import encode_patch +from modal_dojo.train_recipes.miles_recipe.recipe import MilesRecipe + +_TRAIN_DISK_MIB = 2048 * 1024 +_PATCH_DIR = ( + Path(__file__).resolve().parents[2] / "frameworks/miles/modal_helpers/patches" +) + + +def _image_patches() -> list[str]: + return [ + f"echo {encode_patch(name, _PATCH_DIR)} | base64 -d | python3" + for name in ( + "patch_glm_5_3_flash_kda", + "patch_glm_5_3_flash_timing", + "patch_glm_5_3_flash_fp8_device", + ) + ] + + +@dataclass(config=ConfigDict(extra="forbid", arbitrary_types_allowed=True)) +class GLM_5_3_Flash_Recipe(MilesRecipe): + """GLM-5.3-Flash text GRPO on four nodes of eight B300 GPUs.""" + + model_config_class: ClassVar[type[ModelConfig]] = GLM_5_3_Flash + + # glm53next supplies the matching SGLang and Megatron implementations. + docker_image: str = ( + "radixark/miles:glm53next@sha256:" + "66725f740a6013b00d27e21fdfd480a24b0d3b5c61840342e9405bf1e09e5162" + ) + miles_git_ref: str | None = "cc76e23915b2132ecfff65b97331fd83b02637e6" + image_run_commands: list[str] = field(default_factory=_image_patches) + miles_model_name: str = "glm5.3-flash" + model_name: str = "glm5_next" + + gpu_type: str = "B300" + memory: tuple[int, int] = (1024, 2 * 1024 * 1024) + actor_num_nodes: int = 4 + actor_num_gpus_per_node: int = 8 + tensor_model_parallel_size: int = 8 + pipeline_model_parallel_size: int = 4 + expert_model_parallel_size: int = 8 + sequence_parallel: bool = True + decoder_first_pipeline_num_layers: int = 11 + decoder_last_pipeline_num_layers: int = 12 + + environment: dict[str, str] = field( + default_factory=lambda: { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "CONVERT_KEEP_PP1": "1", + "NCCL_MNNVL_ENABLE": "0", + "SGLANG_HEALTH_CHECK_TIMEOUT": "120", + # Inductor's subprocess pool deadlocks under torch_memory_saver. + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + ) + + # The raw converter combines the three mHC alpha parameters on weight sync. + megatron_to_hf_mode: str = "raw" + ref_load: str = "/checkpoints/GLM-5.3-Flash_torch_dist" + conversion_tensor_model_parallel_size: int = 8 + conversion_pipeline_model_parallel_size: int = 1 + conversion_expert_model_parallel_size: int = 8 + conversion_expert_tensor_parallel_size: int = 1 + convert_ephemeral_disk_mb: int | None = 1024 * 1024 + + rollout_batch_size: int = 4 + n_samples_per_prompt: int = 8 + global_batch_size: int = 32 + rollout_temperature: float = 0.8 + rm_type: str | None = "math" + balance_data: bool = True + skip_eval_before_train: bool = True + use_dynamic_batch_size: bool = False + micro_batch_size: int = 1 + max_tokens_per_gpu: int = 8192 + recompute_granularity: str = "full" + recompute_method: str = "uniform" + recompute_num_layers: int = 1 + bf16: bool = True + use_distributed_optimizer: bool = True + + offload_train: bool = True + offload_train_target: str = "disk" + offload_train_disk_dir: str = "/tmp/train_offload" + train_memory_margin_bytes: int = 3 * 1024**3 + update_weight_buffer_size: int = 1024**3 + no_save_optim: bool = True + no_load_optim: bool = True + no_save_rng: bool = True + no_load_rng: bool = True + train_function_kwargs: dict[str, Any] = field( + default_factory=lambda: {"ephemeral_disk": _TRAIN_DISK_MIB} + ) + + rollout_num_gpus_per_engine: int = 8 + sglang_device: str = "cuda" + sglang_tp_size: int = 8 + sglang_ep_size: int = 8 + sglang_moe_runner_backend: str = "triton" + sglang_mem_fraction_static: float = 0.7 + sglang_chunked_prefill_size: int = 8192 + sglang_disable_radix_cache: bool = True + sglang_dsa_prefill_backend: str = "tilelang" + sglang_dsa_decode_backend: str = "tilelang" + sglang_kv_cache_dtype: str = "bfloat16" + check_weight_update_equal: bool = True + check_weight_update_skip_list: list[str] = field( + default_factory=lambda: ["visual."] + ) + router_health_success_threshold: int = 1 + router_health_check_interval_secs: int = 15 + router_health_failure_threshold: int = 40 + rollout_health_check_interval: int = 300 + rollout_health_check_timeout: int = 300 + rollout_health_check_first_wait: int = 1800 + distributed_timeout_minutes: int = 60 + + @model_validator(mode="after") + def _keep_disk_reservation(self) -> "GLM_5_3_Flash_Recipe": + kwargs = self.train_function_kwargs or {} + if "ephemeral_disk" not in kwargs: + object.__setattr__( + self, + "train_function_kwargs", + {"ephemeral_disk": _TRAIN_DISK_MIB, **kwargs}, + ) + return self + + @model_validator(mode="after") + def _keep_image_patches(self) -> "GLM_5_3_Flash_Recipe": + patches = _image_patches() + current = self.image_run_commands or [] + if current[: len(patches)] != patches: + object.__setattr__( + self, + "image_run_commands", + [*patches, *(c for c in current if c not in patches)], + ) + return self diff --git a/modal_dojo/train_recipes/miles_recipe/recipe.py b/modal_dojo/train_recipes/miles_recipe/recipe.py index 15d9370c4..2e891ab18 100644 --- a/modal_dojo/train_recipes/miles_recipe/recipe.py +++ b/modal_dojo/train_recipes/miles_recipe/recipe.py @@ -905,6 +905,9 @@ def get_base_recipe(cls, model_config: ModelConfig) -> "MilesRecipe | None": from modal_dojo.train_recipes.miles_recipe.gemma4_26b_a4b import ( Gemma4_26B_A4B_Recipe, ) + from modal_dojo.train_recipes.miles_recipe.glm_5_3_flash import ( + GLM_5_3_Flash_Recipe, + ) from modal_dojo.common.models.glm_5_3_flash import GLM_5_3_Flash_LoRA from modal_dojo.train_recipes.miles_recipe.glm_5_3_flash_lora import ( GLM_5_3_Flash_LoRA_Recipe, @@ -926,6 +929,8 @@ def get_base_recipe(cls, model_config: ModelConfig) -> "MilesRecipe | None": if isinstance(model_config, GLM_5_3_Flash_LoRA): return GLM_5_3_Flash_LoRA_Recipe() + if model_config.model_name == "zai-org/GLM-5.3-Flash": + return GLM_5_3_Flash_Recipe() if model_config.model_name == "Qwen/Qwen3.5-4B": return Qwen3_5_4B_Miles_Recipe() if model_config.model_name == "moonshotai/Moonlight-16B-A3B-Instruct":