Print step time for each step

[samos123]

This is helpful in cases where there is variable step time and looking at the logs would quickly allow you to identify such cases.

[markblee]

Hm, I'm not sure that we want to print every step by default (since it'll introduce a significant amount of logs/summaries, and since avg step time is usually pretty stable after the first couple hundred steps)?

[samos123]

That makes sense. Should I hide it behind an option? For me it was important to troubleshoot a variable step time issue. Or would you rather totally leave this out of the code base. Note I'm fine with that too.

[samos123]

I've added it as a config parameter and made it false by default

[samos123]

I have been using this for 2 use cases:

• Ensure step time is stable across steps. In the past on the GPU, the networking becomes unstable where collectives sometimes take longer between steps. E.g. step time varies between 5 second to 7 seconds.
• Quickly evaluate whether there is a performance gain without having to wait for 100 steps

Please let me know if there is any interest in merging it @markblee

[markblee]

Maybe it'll be more general to do:

Suggested change

	# Log the step time for each individual step
	log_step_time: bool = False
	# Interval to log average step time. If None, defaults to 100.
	log_step_time_every_n_steps: Optional[int] = None

Note that we default to None instead of 100 to avoid triggering a large number of golden config updates.

[markblee]

Suggested change

	if num_steps % 100 == 0:
	if num_steps % (cfg.log_step_time_every_n_steps or 100) == 0: