MCPcopy Create free account
hub / github.com/AI-Hypercomputer/maxtext / create_engine_from_config_flags

Function create_engine_from_config_flags

src/MaxText/maxengine.py:1670–1706  ·  view source on GitHub ↗

Create new MaxEngine instance with given batch_size, prefill and target lengths, and any config params provided through `args_str`.

(
    maxengine_config_filepath, batch_size, max_prefill_predict_length, max_target_length, args_str
)

Source from the content-addressed store, hash-verified

1668
1669
1670def create_engine_from_config_flags(
1671 maxengine_config_filepath, batch_size, max_prefill_predict_length, max_target_length, args_str
1672):
1673 """Create new MaxEngine instance with given batch_size, prefill and target lengths, and any config
1674 params provided through `args_str`.
1675 """
1676 # batch and cache related
1677 args = {
1678 "scan_layers": "false",
1679 "async_checkpointing": "false",
1680 "ici_fsdp_parallelism": "1",
1681 "ici_autoregressive_parallelism": "1",
1682 "ici_tensor_parallelism": "-1",
1683 "weight_dtype": "bfloat16",
1684 "attention": "dot_product",
1685 "max_prefill_predict_length": f"{max_prefill_predict_length}",
1686 "max_target_length": f"{max_target_length}",
1687 "per_device_batch_size": f"{batch_size}",
1688 }
1689
1690 print(f"Command line args: {args_str}")
1691 cmd_args = args_str.split(" ")
1692 for cmd_arg in cmd_args:
1693 if cmd_arg:
1694 k, v = cmd_arg.split("=")
1695 args[k.strip()] = v.strip()
1696 assert "load_parameters_path" in args, "load_parameters_path must be defined"
1697 if maxengine_config_filepath is None:
1698 maxengine_config_filepath = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")
1699 updated_args = [os.path.join(MAXTEXT_PKG_DIR, "maxengine_server.py"), maxengine_config_filepath]
1700 for k, v in args.items():
1701 option = f"{k}={v}"
1702 updated_args.append(option)
1703 print(f"Invoking maxengine with args:\n \t{updated_args}")
1704 cfg = pyconfig.initialize(updated_args)
1705 engine = MaxEngine(cfg)
1706 return engine

Callers 1

mainFunction · 0.90

Calls 2

MaxEngineClass · 0.85
initializeMethod · 0.45

Tested by

no test coverage detected