Skip to content

Commit

Permalink
simplify
Browse files Browse the repository at this point in the history
  • Loading branch information
irenedea committed Oct 23, 2024
1 parent c681a9c commit 34541f1
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions llmfoundry/command_utils/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -311,18 +311,19 @@ def train(cfg: DictConfig) -> Trainer:
eval_gauntlet_config = train_cfg.eval_gauntlet or train_cfg.eval_gauntlet_str

# Optional parameters will be set to default values if not specified.
env_run_name: Optional[str] = os.environ.get('RUN_NAME', None)
run_name: Optional[str] = (
train_cfg.run_name if train_cfg.run_name else env_run_name
)
run_name: Optional[
str] = train_cfg.run_name if train_cfg.run_name else os.environ.get(
'RUN_NAME',
None,
)
is_state_dict_sharded: bool = (
fsdp_config.get('state_dict_type', 'full') == 'sharded'
) if fsdp_config else False
save_latest_filename: str = train_cfg.save_latest_filename if train_cfg.save_latest_filename else 'latest-sharded-rank{rank}' if is_state_dict_sharded else 'latest-rank{rank}.pt'
save_filename: str = train_cfg.save_filename if train_cfg.save_filename else 'ep{epoch}-ba{batch}-rank{rank}.pt'

# Enable autoresume from model checkpoints if possible
is_user_set_run_name: bool = train_cfg.run_name is not None or env_run_name is not None
is_user_set_run_name: bool = train_cfg.run_name is not None or run_name is not None
autoresume_default: bool = False
if is_user_set_run_name and \
train_cfg.save_folder is not None \
Expand Down

0 comments on commit 34541f1

Please sign in to comment.