32 b #121
32 b #121
Conversation
| @@ -130,7 +130,7 @@ def build(self, trainer: "Trainer") -> Optional[Callback]: | |||
| eval_batch_size = ( | |||
| self.eval_batch_size | |||
| if self.eval_batch_size is not None | |||
| else trainer.rank_microbatch_size * get_world_size(trainer.dp_process_group) | |||
| else 2 * trainer.rank_microbatch_size * get_world_size(trainer.dp_process_group) | |||
There was a problem hiding this comment.
nit: you could instead passed an updated evaluator callback in OLMo2-32B.py:
.with_callback(
"lm_evaluator",
LMEvaluatorCallbackConfig(
eval_batch_size=<whatever you want>,
eval_dataset=NumpyDatasetConfig.from_data_mix(
DataMix.v3_small_ppl_validation,
name=NumpyDatasetType.padded_fsl,
mix_base_dir=root_dir,
sequence_length=dataset_config.effective_sequence_length,
tokenizer=tokenizer_config,
work_dir=get_work_dir(root_dir),
),
eval_interval=1000,
),
There was a problem hiding this comment.
Yeah, but I think this is better. I think we can default to 2x the training batch size. It should always work.
| return ( | ||
| TrainerConfig( | ||
| save_folder=common.save_folder, | ||
| save_folder=f"gs://ai2-llm/checkpoints/{project_name}/", |
There was a problem hiding this comment.
It defaults to something under my name? Not what we want for an official run?
There was a problem hiding this comment.
Especially if we swap babysitting responsibilities during the run
| # import flash_attn.ops.triton.cross_entropy as flash_attn_ce # type: ignore | ||
|
|
||
| _fused_cross_entropy_loss = triton_ce_loss.cross_entropy_loss | ||
| import flash_attn.ops.triton.cross_entropy as flash_attn_ce # type: ignore |
There was a problem hiding this comment.
Our in-house triton CE loss was copied directly from the flash-attn repo, so I don't see the point of this.
There was a problem hiding this comment.
Ok, I took this back out.
There was a problem hiding this comment.
Do I want compiling and fused loss at the same time?
| """ | ||
| d_model = 5120 |
There was a problem hiding this comment.
this is a very narrow model then... are you sure about that?
There was a problem hiding this comment.
It's a clone of Qwen 32. The tradeoffs are, narrow d_model, wide FFN, GQA, lots of layers.
src/scripts/train/OLMo2-32B.py
Outdated
| fused_loss=True, | ||
| compile_loss=False, |
There was a problem hiding this comment.
I understand the trepidation about the different loss implementations, but the way it was before was the most performant. This way will be slower and have a higher memory footprint.
There was a problem hiding this comment.
Can we have some certainty that this will do the right thing? What happens if we take the 13B from a late checkpoint and run it?
| enabled=False, | ||
| cancel_check_interval=10, | ||
| ), | ||
| ).with_callback( |
There was a problem hiding this comment.
We should just add this to the common callbacks.
There was a problem hiding this comment.
I don't know that we want these for everything. Default should probably be only the new, blessed ones.
| @@ -85,10 +94,57 @@ def build_trainer_config(common: CommonComponents) -> TrainerConfig: | |||
| WandBCallback( | |||
| name=common.run_name, | |||
| entity="ai2-llm", | |||
| project="OLMo-core-26B", | |||
| project=project_name, | |||
| enabled=False, | |||
There was a problem hiding this comment.
Intentionally disabled still? Just checking
No description provided.