Decouple the throughput monitor from the dataset

[braceal]

We should be able to set this using trainer max_steps. That would fully decouple it from the specific dataset. Currently, this code depends on the Dataset class implementing small_subset.

# Select size of subset to use, more ranks require more data to compute stats.
nodes_to_sample_size = {1: 1600, 2: 1600, 4: 3200, 8: 6400}
new_config.small_subset = nodes_to_sample_size.get(cfg.num_nodes, 16000)