JetStream Offline Engine

[wenxindongwork]

JetStream Offline Engine for GRPO

This pull request introduces the JetStream Offline Engine offline_engine.py, an inference engine implemented to support sampling in GRPO, but also intended to be used for generic RL sampling and batch inference.

Example Usage:

offline_engine = OfflineEngine(
    config=maxtext_config,
    params=None,
    enable_batch_prefill=True,
)

input_data = [
    jax.numpy.arange(80),
    jax.numpy.arange(90),
    jax.numpy.arange(100),
]

results = offline_engine.batch_inference(input_data)

for completion_output in results:
    text = offline_engine.tokenizer.decode(completion_output.token_ids)
    max_logging.log(f"Output: {text}")

This Offline Engine works on both McJAX and Pathways. Currently it supports continuous batching and batch prefill. Support for multisampling is yet to be added.

Logical Axis Updates for supporting DP

This PR also introduces changes to support a new logical axis for supporting sequence parallel during Prefill. This is required for supporting Offline Engine with our version of data parallel where prefill shards sequence dimension and decode shards batch dimension.

• Added PREFILL_LENGTH to MaxText/common_types.py to define prefill activation lengths.
• Updated logical axis rules in MaxText/configs/base.yml and MaxText/configs/inference.yml to include prefill_activation_length and prefill_activation_norm_length.
• Modified logical constraints in MaxText/layers/deepseek.py, MaxText/layers/llama2.py , MaxText/layers/gemma2.py to dynamically adjust activation sharding based on MODEL_MODE_PREFILL.
• Modified MaxText/layers/attentions.py to apply logical constraints specific to prefill mode for attention weights, masks, and outputs.
• Modified MaxText/layers/embeddings.py to accept model_mode as an argument for embedding operations
• Modified logical constraints in MaxText/layers/linears.py to apply logical constraints specific to prefill mode

RL specific changes

• Added @jax.jit to log_prob_of_chosen_token in MaxText/inference_utils.py for improved performance.
• Added logic in prefill_packing.py to avoid pre-compiling the batch prefill function as it did not work on Pathways.
• Added logic in prefill_packing.py to conditionally return prompt token log probabilities. This is needed for GRPO loss calculation.

Tests

Added offline_engine_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[khatwanimohit]

@vipannalla @mitalisi Can you help review this PR?