global_automodel#

Classes

GlobalKDConfig

GlobalKDResult

KDLossTermConfig

Functions

build_global_kd_config

build_automodel_global_kd_recipe

Translate Puzzletron KD settings into the current NeMo AutoModel recipe schema.

run_automodel_global_kd

run_global_kd

class GlobalKDConfig#

Bases: object

__init__(*, teacher_dir, student_dir, output_dir, descriptor=None, teacher_descriptor=None, student_descriptor=None, force_hf=True, teacher_force_hf=None, student_force_hf=None, teacher_model_kwargs=<factory>, student_model_kwargs=<factory>, domain='auto', trust_remote_code=True, torch_dtype='bf16', attn_implementation=None, tp=1, pp=1, ep=1, cp=1, dp=1, sequence_parallel=False, activation_checkpointing=False, pp_schedule='1f1b', save_consolidated=False, checkpoint_format='auto', main_ce_weight=1.0, mtp_ce_weight=0.0, main_kd_weight=1.0, mtp_kd_weight=0.0, main_kd=<factory>, mtp_kd=<factory>, ce_weight=None, kd_weight=None, hidden_kd_weight=0.0, temperature=None, global_batch_size=128, local_batch_size=1, max_steps=1, checkpoint_every_steps=None, packed_sequence_size=0, lr=1e-05, weight_decay=0.0, seed=1111, dataset_name='rajpurkar/squad', dataset_split='train', validation_enabled=True, validation_split='validation', resume=True, metadata=<factory>, data=<factory>, freeze_policy='vision_frozen')#
Parameters:

  • teacher_dir (Path)

  • student_dir (Path)

  • output_dir (Path)

  • descriptor (str | None)

  • teacher_descriptor (str | None)

  • student_descriptor (str | None)

  • force_hf (bool)

  • teacher_force_hf (bool | None)

  • student_force_hf (bool | None)

  • teacher_model_kwargs (dict[str, Any])

  • student_model_kwargs (dict[str, Any])

  • domain (Literal['auto', 'llm', 'vlm'])

  • trust_remote_code (bool)

  • torch_dtype (str)

  • attn_implementation (str | None)

  • tp (int)

  • pp (int)

  • ep (int)

  • cp (int)

  • dp (int)

  • sequence_parallel (bool)

  • activation_checkpointing (bool | str)

  • pp_schedule (Literal['1f1b', 'interleaved1f1b'])

  • save_consolidated (bool | str)

  • checkpoint_format (Literal['auto', 'safetensors', 'torch_save'])

  • main_ce_weight (float)

  • mtp_ce_weight (float)

  • main_kd_weight (float)

  • mtp_kd_weight (float)

  • main_kd (KDLossTermConfig)

  • mtp_kd (KDLossTermConfig)

  • ce_weight (float | None)

  • kd_weight (float | None)

  • hidden_kd_weight (float)

  • temperature (float | None)

  • global_batch_size (int)

  • local_batch_size (int)

  • max_steps (int)

  • checkpoint_every_steps (int | None)

  • packed_sequence_size (int)

  • lr (float)

  • weight_decay (float)

  • seed (int)

  • dataset_name (str)

  • dataset_split (str)

  • validation_enabled (bool)

  • validation_split (str)

  • resume (bool)

  • metadata (dict[str, Any])

  • data (dict[str, Any])

  • freeze_policy (Literal['vision_frozen', 'projector_and_language', 'train_all'])

Return type:

None

activation_checkpointing: bool | str = False#
attn_implementation: str | None = None#
ce_weight: float | None = None#
checkpoint_every_steps: int | None = None#
checkpoint_format: Literal['auto', 'safetensors', 'torch_save'] = 'auto'#
cp: int = 1#
data: dict[str, Any]#
dataset_name: str = 'rajpurkar/squad'#
dataset_split: str = 'train'#
descriptor: str | None = None#
domain: Literal['auto', 'llm', 'vlm'] = 'auto'#
dp: int = 1#
ep: int = 1#
force_hf: bool = True#
freeze_policy: Literal['vision_frozen', 'projector_and_language', 'train_all'] = 'vision_frozen'#
global_batch_size: int = 128#
hidden_kd_weight: float = 0.0#
property identity: str#
kd_weight: float | None = None#
local_batch_size: int = 1#
lr: float = 1e-05#
main_ce_weight: float = 1.0#
main_kd: KDLossTermConfig#
main_kd_weight: float = 1.0#
max_steps: int = 1#
metadata: dict[str, Any]#
mtp_ce_weight: float = 0.0#
mtp_kd: KDLossTermConfig#
mtp_kd_weight: float = 0.0#
property needs_teacher: bool#
property objective_weights: dict[str, float]#
output_dir: Path#
packed_sequence_size: int = 0#
pp: int = 1#
pp_schedule: Literal['1f1b', 'interleaved1f1b'] = '1f1b'#
property resolved_student_descriptor: str | None#
property resolved_student_force_hf: bool#
property resolved_teacher_descriptor: str | None#
property resolved_teacher_force_hf: bool#
resume: bool = True#
save_consolidated: bool | str = False#
seed: int = 1111#
sequence_parallel: bool = False#
student_descriptor: str | None = None#
student_dir: Path#
student_force_hf: bool | None = None#
student_model_kwargs: dict[str, Any]#
teacher_descriptor: str | None = None#
teacher_dir: Path#
teacher_force_hf: bool | None = None#
teacher_model_kwargs: dict[str, Any]#
temperature: float | None = None#
to_dict()#
Return type:

dict[str, Any]

torch_dtype: str = 'bf16'#
tp: int = 1#
trust_remote_code: bool = True#
validation_enabled: bool = True#
validation_split: str = 'validation'#
weight_decay: float = 0.0#
class GlobalKDResult#

Bases: object

__init__(*, kd_id, output_dir, metrics=<factory>)#
Parameters:

  • kd_id (str)

  • output_dir (Path)

  • metrics (dict[str, Any])

Return type:

None

kd_id: str#
metrics: dict[str, Any]#
output_dir: Path#
to_dict()#
Return type:

dict[str, Any]

class KDLossTermConfig#

Bases: object

__init__(*, metric='kld', temperature=1.0, chunk_size=0)#
Parameters:

  • metric (Literal['kld', 'tvd'])

  • temperature (float)

  • chunk_size (int)

Return type:

None

chunk_size: int = 0#
metric: Literal['kld', 'tvd'] = 'kld'#
temperature: float = 1.0#
build_automodel_global_kd_recipe(kd_config)#

Translate Puzzletron KD settings into the current NeMo AutoModel recipe schema.

Parameters:

kd_config (GlobalKDConfig)

Return type:

dict[str, Any]

build_global_kd_config(config)#
Parameters:

config (dict[str, Any])

Return type:

GlobalKDConfig

run_automodel_global_kd(kd_config)#
Parameters:

kd_config (GlobalKDConfig)

Return type:

dict[str, Any]

run_global_kd(kd_config, recipe_runner=None)#
Parameters:

kd_config (GlobalKDConfig)

Return type:

GlobalKDResult