Llama

Eden11BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~14B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden11BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~14B (keeps all Eden behaviors)."""

    # If you want long context like Eden-long, bump this; else inherit 8192.
    seq_length: int = 8192  # or remove this line to keep 8192

    # ~14B sizing (head_dim ≈ 128)
    num_layers: int = 36
    hidden_size: int = 5120
    ffn_hidden_size: int = 13824
    num_attention_heads: int = 40
    num_query_groups: int = 8  # GQA (inherited value is also fine if already 8)

Eden18BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~18B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden18BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~18B (keeps all Eden behaviors)."""

    # If you want long context like Eden-long, bump this; else inherit 8192.
    seq_length: int = 8192  # or remove this line to keep 8192

    # ~18B sizing (head_dim ≈ 128)
    num_layers: int = 48
    hidden_size: int = 6144
    ffn_hidden_size: int = 16384
    num_attention_heads: int = 48
    num_query_groups: int = 8  # GQA (inherited value is also fine if already 8)
    old_context_len: int = 8192  # or remove this line to keep 8192

Eden21BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~21B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden21BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~21B (keeps all Eden behaviors)."""

    seq_length: int = 8192

    # ~21B sizing (head_dim = 128)
    num_layers: int = 42  # 42 layers for 21B target
    hidden_size: int = 7168  # 56 * 128 = 7168 for exact head_dim
    ffn_hidden_size: int = 19456  # ~2.7x hidden_size
    num_attention_heads: int = 56  # Divisible by 8
    num_query_groups: int = 8  # GQA
    old_context_len: int = 8192

Eden24BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~8B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden24BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~8B (keeps all Eden behaviors)."""

    # If you want long context like Eden-long, bump this; else inherit 8192.
    seq_length: int = 32768  # or remove this line to keep 8192

    # ~8B sizing (head_dim ≈ 128)
    num_layers: int = 46
    hidden_size: int = 6144
    ffn_hidden_size: int = 23296
    num_attention_heads: int = 48
    num_query_groups: int = 8  # GQA (inherited value is also fine if already 8)
    old_context_len: int = 8192

Eden27BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~8B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden27BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~8B (keeps all Eden behaviors)."""

    # If you want long context like Eden-long, bump this; else inherit 8192.
    seq_length: int = 32768  # or remove this line to keep 8192

    # ~8B sizing (head_dim ≈ 128)
    num_layers: int = 46
    hidden_size: int = 6656
    ffn_hidden_size: int = 23296
    num_attention_heads: int = 52
    num_query_groups: int = 8  # GQA (inherited value is also fine if already 8)
    old_context_len: int = 8192

Eden28BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~28B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden28BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~28B (keeps all Eden behaviors)."""

    # If you want long context like Eden-long, bump this; else inherit 8192.
    seq_length: int = 8192  # or remove this line to keep 8192

    # ~8B sizing (head_dim ≈ 128)
    num_layers: int = 48
    hidden_size: int = 6144
    ffn_hidden_size: int = 26368
    num_attention_heads: int = 48
    num_query_groups: int = 8  # GQA (inherited value is also fine if already 8)
    old_context_len: int = 8192  # or remove this line to keep 8192

Eden35BConfig dataclass

Bases: EdenConfig

Eden-flavoured Llama-3.1 ~35B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class Eden35BConfig(EdenConfig):
    """Eden-flavoured Llama-3.1 ~35B (keeps all Eden behaviors)."""

    seq_length: int = 8192

    # ~35B sizing (head_dim ≈ 128)
    num_layers: int = 64
    hidden_size: int = 7168
    ffn_hidden_size: int = 20480
    num_attention_heads: int = 56
    num_query_groups: int = 8  # GQA
    old_context_len: int = 8192

EdenConfig dataclass

Bases: Llama3Config8B

Eden-flavoured Llama-3.1 ~8B (keeps all Eden behaviors).

Source code in bionemo/evo2/models/llama.py

@dataclass
class EdenConfig(llm.Llama3Config8B):
    """Eden-flavoured Llama-3.1 ~8B (keeps all Eden behaviors)."""

    rotary_base: int = 500_000
    seq_length: int = 8192
    num_layers: int = 32
    hidden_size: int = 4096
    ffn_hidden_size: int = 14336
    num_attention_heads: int = 32

    scale_factor: int = 1
    low_freq_factor: int = 1
    high_freq_factor: int = 4
    old_context_len: int = 8192
    init_method_std: float = 0.02
    embedding_init_method_std: Optional[float] = None

    def configure_model(self, *args, **kwargs):
        """Configure and instantiate a Megatron Core Llama 3.1 model.

        Extends the base configuration with Llama 3.1 specific RoPE scaling.
        """
        model = super(EdenConfig, self).configure_model(*args, **kwargs)
        # Apply rope scaling for Llama3.1 model
        model.rotary_pos_emb.inv_freq = apply_rope_scaling(
            model.rotary_pos_emb.inv_freq,
            factor=self.scale_factor,
            low_freq_factor=self.low_freq_factor,
            high_freq_factor=self.high_freq_factor,
            old_context_len=self.old_context_len,
        )
        return model

configure_model(*args, **kwargs)

Configure and instantiate a Megatron Core Llama 3.1 model.

Extends the base configuration with Llama 3.1 specific RoPE scaling.

Source code in bionemo/evo2/models/llama.py

def configure_model(self, *args, **kwargs):
    """Configure and instantiate a Megatron Core Llama 3.1 model.

    Extends the base configuration with Llama 3.1 specific RoPE scaling.
    """
    model = super(EdenConfig, self).configure_model(*args, **kwargs)
    # Apply rope scaling for Llama3.1 model
    model.rotary_pos_emb.inv_freq = apply_rope_scaling(
        model.rotary_pos_emb.inv_freq,
        factor=self.scale_factor,
        low_freq_factor=self.low_freq_factor,
        high_freq_factor=self.high_freq_factor,
        old_context_len=self.old_context_len,
    )
    return model

HFEdenLlamaImporter

Bases: HFLlamaImporter

Importer for Eden-flavoured Llama models which just overrides the tokenizer and config classes from NeMo.

Source code in bionemo/evo2/models/llama.py

@io.model_importer(LlamaModel, "hf")
class HFEdenLlamaImporter(HFLlamaImporter):
    """Importer for Eden-flavoured Llama models which just overrides the tokenizer and config classes from NeMo."""

    @property
    def config(self) -> EdenConfig:
        """Create a NeMo LlamaConfig from the HF model config.

        Translates the HF configuration parameters to the equivalent NeMo
        configuration.

        Returns:
            LlamaConfig: NeMo configuration for Llama models
        """
        from transformers import AutoConfig, GenerationConfig

        source = AutoConfig.from_pretrained(str(self))
        try:
            generation_config = GenerationConfig.from_pretrained(str(self))
        except Exception:
            generation_config = None

        def make_vocab_size_divisible_by(vocab_size):
            base = 128
            while vocab_size % base != 0:
                base //= 2
            return base

        cls = EdenConfig
        scale_factor = source.rope_scaling.get("factor", 8.0) if source.rope_scaling is not None else 8.0

        args = {}

        output = cls(
            num_layers=source.num_hidden_layers,
            hidden_size=source.hidden_size,
            ffn_hidden_size=(
                source.intermediate_size
                if not getattr(source, "intermediate_size_mlp", None)
                else source.intermediate_size_mlp
            ),
            num_attention_heads=source.num_attention_heads,
            init_method_std=source.initializer_range,
            layernorm_epsilon=source.rms_norm_eps,
            num_query_groups=source.num_key_value_heads,
            seq_length=source.max_position_embeddings,
            rotary_base=source.rope_theta,
            gated_linear_unit=True,
            make_vocab_size_divisible_by=make_vocab_size_divisible_by(source.vocab_size),
            share_embeddings_and_output_weights=getattr(source, "tie_word_embeddings", False),
            fp16=(dtype_from_hf(source) == torch.float16),
            bf16=(dtype_from_hf(source) == torch.bfloat16),
            params_dtype=dtype_from_hf(source),
            generation_config=generation_config,
            vocab_size=source.vocab_size,
            kv_channels=getattr(source, "head_dim", None),
            scale_factor=scale_factor,
            **args,
        )

        return output

    @property
    def tokenizer(self):
        """Override the tokenizer to use the Eden-flavoured tokenizer."""
        from bionemo.evo2.run.utils import patch_eden_tokenizer  # avoid circular import

        tokenizer = get_nmt_tokenizer("byte-level")
        patch_eden_tokenizer(tokenizer)
        return tokenizer

config property

Create a NeMo LlamaConfig from the HF model config.

Translates the HF configuration parameters to the equivalent NeMo configuration.

Returns:

Name	Type	Description
LlamaConfig	EdenConfig	NeMo configuration for Llama models

tokenizer property

Override the tokenizer to use the Eden-flavoured tokenizer.