Initialize the OfflineEngine. Args: config: The MaxText config object which will be used to create MaxEngine instance(s). params: Model parameters (loaded from engine if None) enable_batch_prefill: Whether to use prefill packing. config.scan_layers
(
self,
config: Any,
params: None | Params = None,
enable_batch_prefill: bool = False,
min_decode_steps: int = 10,
tokenizer: Any = None,
eos_ids: list[int] | None = None,
prefill_lengths: list[int] | str = "auto",
batch_prefill_max_batch_size: int = 16,
mesh: Mesh = None,
rng: jax.random.PRNGKey = None,
debug: bool = False,
)
| 755 | """Class for handling offline inference on batches of inputs.""" |
| 756 | |
| 757 | def __init__( |
| 758 | self, |
| 759 | config: Any, |
| 760 | params: None | Params = None, |
| 761 | enable_batch_prefill: bool = False, |
| 762 | min_decode_steps: int = 10, |
| 763 | tokenizer: Any = None, |
| 764 | eos_ids: list[int] | None = None, |
| 765 | prefill_lengths: list[int] | str = "auto", |
| 766 | batch_prefill_max_batch_size: int = 16, |
| 767 | mesh: Mesh = None, |
| 768 | rng: jax.random.PRNGKey = None, |
| 769 | debug: bool = False, |
| 770 | ): |
| 771 | """Initialize the OfflineEngine. |
| 772 | |
| 773 | Args: |
| 774 | config: The MaxText config object which will be used to |
| 775 | create MaxEngine instance(s). |
| 776 | params: Model parameters (loaded from engine if None) |
| 777 | enable_batch_prefill: Whether to use prefill packing. |
| 778 | config.scan_layers must be False if this is True |
| 779 | min_decode_steps: Number of decode steps to perform at a time, |
| 780 | before checking for completion. |
| 781 | eos_ids: list of EOS token IDs for checking sequence completion. |
| 782 | If None, the tokenizer's EOS token will be used. |
| 783 | tokenizer: Tokenizer instance for encoding/decoding text. If None, |
| 784 | will be created using the config if eos_ids is not provided. |
| 785 | prefill_lengths: list of expected prefill lengths, or "auto" to |
| 786 | automatically determine appropriate lengths from the engine |
| 787 | config. Input sequences will be padded to the nearest length |
| 788 | in this list. |
| 789 | batch_prefill_max_batch_size: Maximum number of inputs to pack |
| 790 | into a single prefill. This is only used when enable_batch_prefill |
| 791 | is True. |
| 792 | mesh: JAX Mesh object. Use this |
| 793 | argument if you want to use only some of the devices for OfflineEngine and |
| 794 | reserve the rest for other tasks. If None, OfflineEngine will create the mesh |
| 795 | automatically. |
| 796 | rng: Random number generator key. If None, a new key will be created. |
| 797 | """ |
| 798 | max_logging.log("Initializing OfflineEngine") |
| 799 | # Configurations |
| 800 | self.config = config |
| 801 | self.params = params |
| 802 | self.min_decode_steps = min_decode_steps |
| 803 | self.enable_batch_prefill = enable_batch_prefill |
| 804 | self.mesh = mesh |
| 805 | self.tokenizer = tokenizer |
| 806 | self.eos_ids = eos_ids |
| 807 | self.prefill_lengths = prefill_lengths |
| 808 | self.batch_prefill_max_batch_size = batch_prefill_max_batch_size |
| 809 | self.max_prefill_length = self.config.max_prefill_predict_length |
| 810 | self.max_decode_length = self.config.max_target_length - self.max_prefill_length |
| 811 | self.rng = jax.random.PRNGKey(0) if rng is None else rng |
| 812 | self.debug = debug |
| 813 | self._validate_config() |
| 814 |
nothing calls this directly
no test coverage detected