Bare bones GenerationModule
#324
Conversation
| max_sequence_length: Optional[int] = None, | ||
| rank_microbatch_size: Optional[int] = None, |
There was a problem hiding this comment.
These actually are optional and if provided are used to warm up the RoPE cache (max_seq_len) and an MoE component (rank_microbatch_size).
There was a problem hiding this comment.
Nit: maybe rename this generation_module.py to be more specific. Also consider moving the transformer implementation to its own submodule.
src/olmo_core/generate/generation.py
Outdated
| from olmo_core.generate.config import GenerationConfig | ||
| from olmo_core.generate.selection import temperature_sampling |
There was a problem hiding this comment.
Nit: relative imports are nice to have at least within the same submodule.
| from olmo_core.generate.config import GenerationConfig | |
| from olmo_core.generate.selection import temperature_sampling | |
| from .config import GenerationConfig | |
| from .selection import temperature_sampling |
| Args: | ||
| checkpoint_dir: Path to checkpoint directory | ||
| work_dir: Working directory for caching remote checkpoints | ||
| process_group: Process group for distributed loading | ||
| pre_download: Whether to pre-download remote checkpoints | ||
| load_thread_count: Number of threads to use for loading the checkpoint | ||
|
|
||
| Raises: | ||
| FileNotFoundError: If checkpoint directory doesn't exist | ||
| RuntimeError: If checkpoint loading fails |
There was a problem hiding this comment.
nit: this docstring syntax is inconsistent with our other docstrings
| work_dir = Path( | ||
| work_dir or (tempfile.mkdtemp() if get_rank(process_group) == 0 else "/tmp") | ||
| ) |
There was a problem hiding this comment.
This assumes all ranks share the filesystem? Which is usually only true for single-node jobs.
There was a problem hiding this comment.
Alternatively just force the user to provide a work dir.
Add a simple GenerationModule (analogous to
TrainModule) that can be used to configure and run autoregressive next token prediction.As of now this supports:
1. loading distributed checkpoints exactly as they were saved during training.
2.
temperatureparameter for generation.3. using FSDP to shard larger models across multiple devices (mostly as a demonstration of how other types of parallelism can be worked in).
4. Attention masks passed through to SDPA so that batched generation w/ left-padding is supported.
Note that this implementation is very inefficient compared to
transformersorvllm, in part due to the lack of kv caching.