vtava's picture
Upload verified PDelta3-CLVR checkpoint for layers [3, 7, 11]
a38f163 verified
Raw History Blame
5.66 kB
from .live_console import configure_live_console
# Configure the current process before importing model modules. Every train_*.py
# script imports tinycenn_lm, so notebook-launched trainers inherit immediate,
# line-buffered stdout/stderr even when the notebook uses subprocess.run(...).
configure_live_console()
from .cenn import CeNNConfig, FastCeNNCore
from .modeling import (
DEFAULT_BASE_MODEL,
HybridDecoderLayer,
build_from_adapter,
freeze_for_adapter_training,
inject_cenn,
load_adapter,
save_adapter,
trainable_parameter_summary,
)
from .student import (
CeNNReplacementLayer,
build_cenn_student,
freeze_student_interfaces,
load_cenn_student_weights,
replace_transformer_with_cenn,
save_cenn_student,
student_parameter_summary,
)
from .moe import (
FastMoECeNNCore,
MoECeNNConfig,
MoECeNNReplacementLayer,
build_moe_cenn_student,
freeze_moe_student_interfaces,
load_moe_cenn_student_weights,
moe_router_stats,
replace_transformer_with_moe_cenn,
save_moe_cenn_student,
warmstart_moe_from_plain_cenn,
)
from .sharded_moe import (
FastShardedMoECeNNCore,
ShardedMoECeNNConfig,
ShardedMoECeNNReplacementLayer,
build_sharded_moe_student,
freeze_sharded_moe_interfaces,
load_sharded_moe_student_weights,
replace_transformer_with_sharded_moe_cenn,
save_sharded_moe_student,
sharded_router_stats,
warmstart_sharded_moe_from_plain_cenn,
)
from .story_v2 import (
CausalStoryMemory,
LowRankLMHeadAdapter,
StoryV2Config,
StoryV2ReplacementLayer,
build_story_v2_from_story_v1,
build_story_v2_student,
freeze_story_v2_interfaces,
load_story_v2_weights,
save_story_v2_student,
story_v2_parameter_summary,
story_v2_router_stats,
upgrade_sharded_model_to_story_v2,
)
from .smollm2_amcenn import (
DEFAULT_SMOLLM2,
AMCeNNAttention,
PositiveSoftmaxFeatures,
ShardedTop2LlamaMLP,
SmolAMCeNNConfig,
amcenn_parameter_summary,
amcenn_router_stats,
build_smollm2_amcenn,
freeze_smollm2_for_amcenn_training,
load_smollm2_amcenn_weights,
replace_smollm2_core,
save_smollm2_amcenn,
)
from .smollm2_amcenn_v2 import (
AMCeNNAttentionV2,
AdaptivePositiveSoftmaxFeatures,
SmolAMCeNNV2Config,
build_smollm2_amcenn_v2,
convert_all_ffns_to_sharded_top2,
freeze_for_global_training,
freeze_for_group_calibration,
load_smollm2_amcenn_v2_weights,
replace_all_smollm2_attention,
replace_attention_layers,
save_smollm2_amcenn_v2,
v2_parameter_summary,
)
from .hf_persistence import (
build_model_card,
collect_reports,
install_colab_hf_upload_enhancer,
persist_hf_run,
redact_secrets,
utc_run_id,
)
from .colab_live_backup import (
install_colab_training_backup,
is_tinycenn_training_command,
output_dir_from_command,
)
from .direct_colab_backup import install_direct_training_backup
# In Colab, make Hugging Face backup mandatory. The parent notebook wrapper handles
# normal subprocess-launched trainers. A trainer-side fallback covers notebooks that
# launch train_*.py before the notebook kernel imports tinycenn_lm.
install_colab_training_backup()
install_direct_training_backup()
install_colab_hf_upload_enhancer()
__all__ = [
"configure_live_console",
"CeNNConfig", "FastCeNNCore", "DEFAULT_BASE_MODEL", "HybridDecoderLayer",
"build_from_adapter", "freeze_for_adapter_training", "inject_cenn", "load_adapter",
"save_adapter", "trainable_parameter_summary", "CeNNReplacementLayer", "build_cenn_student",
"freeze_student_interfaces", "load_cenn_student_weights", "replace_transformer_with_cenn",
"save_cenn_student", "student_parameter_summary", "MoECeNNConfig", "FastMoECeNNCore",
"MoECeNNReplacementLayer", "build_moe_cenn_student", "freeze_moe_student_interfaces",
"load_moe_cenn_student_weights", "moe_router_stats", "replace_transformer_with_moe_cenn",
"save_moe_cenn_student", "warmstart_moe_from_plain_cenn", "ShardedMoECeNNConfig",
"FastShardedMoECeNNCore", "ShardedMoECeNNReplacementLayer", "build_sharded_moe_student",
"freeze_sharded_moe_interfaces", "load_sharded_moe_student_weights",
"replace_transformer_with_sharded_moe_cenn", "save_sharded_moe_student", "sharded_router_stats",
"warmstart_sharded_moe_from_plain_cenn", "StoryV2Config", "CausalStoryMemory",
"StoryV2ReplacementLayer", "LowRankLMHeadAdapter", "upgrade_sharded_model_to_story_v2",
"freeze_story_v2_interfaces", "story_v2_router_stats", "story_v2_parameter_summary",
"save_story_v2_student", "load_story_v2_weights", "build_story_v2_student",
"build_story_v2_from_story_v1", "DEFAULT_SMOLLM2", "SmolAMCeNNConfig",
"PositiveSoftmaxFeatures", "AMCeNNAttention", "ShardedTop2LlamaMLP", "replace_smollm2_core",
"freeze_smollm2_for_amcenn_training", "amcenn_router_stats", "amcenn_parameter_summary",
"save_smollm2_amcenn", "load_smollm2_amcenn_weights", "build_smollm2_amcenn",
"SmolAMCeNNV2Config", "AdaptivePositiveSoftmaxFeatures", "AMCeNNAttentionV2",
"convert_all_ffns_to_sharded_top2", "replace_attention_layers", "replace_all_smollm2_attention",
"freeze_for_group_calibration", "freeze_for_global_training", "v2_parameter_summary",
"save_smollm2_amcenn_v2", "load_smollm2_amcenn_v2_weights", "build_smollm2_amcenn_v2",
"build_model_card", "collect_reports", "persist_hf_run", "install_colab_hf_upload_enhancer",
"redact_secrets", "utc_run_id", "install_colab_training_backup", "install_direct_training_backup",
"is_tinycenn_training_command", "output_dir_from_command",
]