Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds support for knowledge distillation #380

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+230 −54
Open

Adds support for knowledge distillation #380

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d025471
feat: enable knowledge distillation
RobotSail Dec 24, 2024
c0b9cf3
fix: make linter happy
RobotSail Dec 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions src/instructlab/training/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ class DeepSpeedOffloadStrategy(Enum):

# public API
class DistributedBackend(Enum):
FSDP: str = "fsdp"
DEEPSPEED: str = "deepspeed"
FSDP = "fsdp"
DEEPSPEED = "deepspeed"


# public API
Expand Down Expand Up @@ -121,6 +121,17 @@ class DeepSpeedOptions(BaseModel):
save_samples: int | None = None


# public API
class DistillationConfig(BaseModel):
"""
Config to use when performing knowledge distillation during training.
"""

temperature: float = Field(1.0, gt=0.0)
alpha: float = Field(1.0, le=1.0, ge=0.0)
teacher_path: str


# public API
class ShardingStrategies(Enum):
FULL_SHARD = "FULL_SHARD"
Expand Down Expand Up @@ -179,6 +190,11 @@ class TrainingArgs(BaseModel):
is_padding_free: bool = False # TODO: deprecate
checkpoint_at_epoch: bool = True
accelerate_full_state_at_epoch: bool = True
weight_decay: float = Field(0.0, ge=0.0)

# settings for knowledge distillation
distillation_options: Optional[DistillationConfig] = None
use_distillation: bool = False

mock_data: Optional[bool] = False
mock_data_len: int = 0
Expand Down
Loading
Loading