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
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
4f495bf
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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if possible would love to standardize on using pathlib.Path rather than str paths.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JamesKunstle I see your point, would it make sense for it to be a path when it can also take on a HF reference? I understand that references can technically still be paths, but to a consumer reading it might sound like only local models are accepted. Would str | Path be satisfactory?



# 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: DistillationConfig | None = None
use_distillation: bool = False

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