fix: Nemotron chunk size alias (#20458)
This commit is contained in:
@@ -32,6 +32,7 @@ MLP = "-"
|
|||||||
MOE = "E"
|
MOE = "E"
|
||||||
DEFAULT_LAYERS_BLOCK_TYPE = ["mamba", "moe", "attention", "moe"]
|
DEFAULT_LAYERS_BLOCK_TYPE = ["mamba", "moe", "attention", "moe"]
|
||||||
DEFAULT_MTP_LAYERS_BLOCK_TYPE = ["attention", "moe"]
|
DEFAULT_MTP_LAYERS_BLOCK_TYPE = ["attention", "moe"]
|
||||||
|
DEFAULT_MAMBA_CHUNK_SIZE = 256
|
||||||
|
|
||||||
|
|
||||||
class NemotronHConfig(PretrainedConfig):
|
class NemotronHConfig(PretrainedConfig):
|
||||||
@@ -213,6 +214,28 @@ class NemotronHConfig(PretrainedConfig):
|
|||||||
mtp_layers_block_type = NemotronHConfig._pattern_to_list(pattern)
|
mtp_layers_block_type = NemotronHConfig._pattern_to_list(pattern)
|
||||||
return mtp_layers_block_type
|
return mtp_layers_block_type
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _resolve_mamba_chunk_size(mamba_chunk_size, kwargs) -> int:
|
||||||
|
"""Resolve canonical mamba_chunk_size from new and legacy config fields."""
|
||||||
|
chunk_size = kwargs.pop("chunk_size", None)
|
||||||
|
if (
|
||||||
|
mamba_chunk_size is not None
|
||||||
|
and chunk_size is not None
|
||||||
|
and mamba_chunk_size != chunk_size
|
||||||
|
):
|
||||||
|
logger.warning(
|
||||||
|
"Both chunk_size=%s and mamba_chunk_size=%s were provided. "
|
||||||
|
"Using mamba_chunk_size.",
|
||||||
|
chunk_size,
|
||||||
|
mamba_chunk_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
if mamba_chunk_size is None:
|
||||||
|
mamba_chunk_size = chunk_size
|
||||||
|
if mamba_chunk_size is None:
|
||||||
|
mamba_chunk_size = DEFAULT_MAMBA_CHUNK_SIZE
|
||||||
|
return mamba_chunk_size
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
vocab_size=131072,
|
vocab_size=131072,
|
||||||
@@ -255,7 +278,7 @@ class NemotronHConfig(PretrainedConfig):
|
|||||||
mamba_dt_init_floor=1e-4,
|
mamba_dt_init_floor=1e-4,
|
||||||
mamba_conv_bias=True,
|
mamba_conv_bias=True,
|
||||||
mamba_proj_bias=False,
|
mamba_proj_bias=False,
|
||||||
mamba_chunk_size=256,
|
mamba_chunk_size=None,
|
||||||
rescale_prenorm_residual=True,
|
rescale_prenorm_residual=True,
|
||||||
n_routed_experts=8,
|
n_routed_experts=8,
|
||||||
n_shared_experts=1,
|
n_shared_experts=1,
|
||||||
@@ -271,6 +294,8 @@ class NemotronHConfig(PretrainedConfig):
|
|||||||
mtp_layers_block_type=DEFAULT_MTP_LAYERS_BLOCK_TYPE,
|
mtp_layers_block_type=DEFAULT_MTP_LAYERS_BLOCK_TYPE,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
):
|
):
|
||||||
|
mamba_chunk_size = self._resolve_mamba_chunk_size(mamba_chunk_size, kwargs)
|
||||||
|
|
||||||
# Compatibility parsing: normalize legacy pattern fields into canonical list fields.
|
# Compatibility parsing: normalize legacy pattern fields into canonical list fields.
|
||||||
layers_block_type = self._resolve_layers_block_type(
|
layers_block_type = self._resolve_layers_block_type(
|
||||||
layers_block_type, hybrid_override_pattern, kwargs
|
layers_block_type, hybrid_override_pattern, kwargs
|
||||||
|
|||||||
Reference in New Issue
Block a user