Adjust default value of max_queue for TF SummaryWriter

GitOrigin-RevId: f37f63f7ca2f685530b93fa93da888cf5fc671a6

Author: Jun Tian

axlearn/common/summary_writer.py

@@ -232,14 +232,15 @@ class Config(BaseWriter.Config):
                 If a type is not listed, `write_every_n_steps` is used as fallback. Each value must
                 be a positive integer multiple of `write_every_n_steps`.
             max_queue: Configures maximum number of summaries before flush.
-                If None, uses the `tf_summary` default (10).
+                Defaults to 1000. The original `tf_summary` default is 10, which is too small
+                and may cause frequent flushes to GCS, potentially blocking training.
             flush_ms: Largest interval between flushes in milliseconds.
                 If None, uses the `tf_summary` default (120,000, i.e. 2 minutes).
         """
 
         write_every_n_steps: int = 1
         write_every_n_steps_map: Optional[dict[SummaryKind, int]] = None
-        max_queue: Optional[int] = None
+        max_queue: int = 1000
         flush_ms: Optional[float] = None
 
     def __init__(self, cfg: BaseWriter.Config, *, parent: Optional[Module]):

axlearn/experiments/testdata/axlearn.experiments.audio.conformer.librispeech_trainer/conformer-l-rnnt.txt

@@ -35,6 +35,7 @@ evalers['eval_train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['eval_train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['eval_train'].metric_calculator.model_method: 'forward'
 evalers['eval_train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['eval_train'].summary_writer.max_queue: 1000
 evalers['eval_train'].summary_writer.write_every_n_steps: 1
 evalers['eval_dev_clean'].eval_dtype: 'jax.numpy.float32'
 evalers['eval_dev_clean'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -63,6 +64,7 @@ evalers['eval_dev_clean'].metric_calculator.klass: 'axlearn.common.evaler.ModelS
 evalers['eval_dev_clean'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['eval_dev_clean'].metric_calculator.model_method: 'forward'
 evalers['eval_dev_clean'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['eval_dev_clean'].summary_writer.max_queue: 1000
 evalers['eval_dev_clean'].summary_writer.write_every_n_steps: 1
 evalers['eval_dev_other'].eval_dtype: 'jax.numpy.float32'
 evalers['eval_dev_other'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -91,6 +93,7 @@ evalers['eval_dev_other'].metric_calculator.klass: 'axlearn.common.evaler.ModelS
 evalers['eval_dev_other'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['eval_dev_other'].metric_calculator.model_method: 'forward'
 evalers['eval_dev_other'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['eval_dev_other'].summary_writer.max_queue: 1000
 evalers['eval_dev_other'].summary_writer.write_every_n_steps: 1
 evalers['decoder_dev_clean'].eval_dtype: 'jax.numpy.float32'
 evalers['decoder_dev_clean'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -128,6 +131,7 @@ evalers['decoder_dev_clean'].metric_calculator.vocab.reverse_extra_ids: True
 evalers['decoder_dev_clean'].metric_calculator.vocab.sentencepiece_model_file: '$DATA_DIR/tokenizers/sentencepiece/librispeech_bpe_1024.model'
 evalers['decoder_dev_clean'].metric_calculator.vocab.use_fast_tokenizer: False
 evalers['decoder_dev_clean'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['decoder_dev_clean'].summary_writer.max_queue: 1000
 evalers['decoder_dev_clean'].summary_writer.write_every_n_steps: 1
 evalers['decoder_dev_other'].eval_dtype: 'jax.numpy.float32'
 evalers['decoder_dev_other'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -165,6 +169,7 @@ evalers['decoder_dev_other'].metric_calculator.vocab.reverse_extra_ids: True
 evalers['decoder_dev_other'].metric_calculator.vocab.sentencepiece_model_file: '$DATA_DIR/tokenizers/sentencepiece/librispeech_bpe_1024.model'
 evalers['decoder_dev_other'].metric_calculator.vocab.use_fast_tokenizer: False
 evalers['decoder_dev_other'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['decoder_dev_other'].summary_writer.max_queue: 1000
 evalers['decoder_dev_other'].summary_writer.write_every_n_steps: 1
 evalers['decoder_test_clean'].eval_dtype: 'jax.numpy.float32'
 evalers['decoder_test_clean'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -202,6 +207,7 @@ evalers['decoder_test_clean'].metric_calculator.vocab.reverse_extra_ids: True
 evalers['decoder_test_clean'].metric_calculator.vocab.sentencepiece_model_file: '$DATA_DIR/tokenizers/sentencepiece/librispeech_bpe_1024.model'
 evalers['decoder_test_clean'].metric_calculator.vocab.use_fast_tokenizer: False
 evalers['decoder_test_clean'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['decoder_test_clean'].summary_writer.max_queue: 1000
 evalers['decoder_test_clean'].summary_writer.write_every_n_steps: 1
 evalers['decoder_test_other'].eval_dtype: 'jax.numpy.float32'
 evalers['decoder_test_other'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -239,6 +245,7 @@ evalers['decoder_test_other'].metric_calculator.vocab.reverse_extra_ids: True
 evalers['decoder_test_other'].metric_calculator.vocab.sentencepiece_model_file: '$DATA_DIR/tokenizers/sentencepiece/librispeech_bpe_1024.model'
 evalers['decoder_test_other'].metric_calculator.vocab.use_fast_tokenizer: False
 evalers['decoder_test_other'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['decoder_test_other'].summary_writer.max_queue: 1000
 evalers['decoder_test_other'].summary_writer.write_every_n_steps: 1
 evalers['decoder_train'].eval_dtype: 'jax.numpy.float32'
 evalers['decoder_train'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -276,6 +283,7 @@ evalers['decoder_train'].metric_calculator.vocab.reverse_extra_ids: True
 evalers['decoder_train'].metric_calculator.vocab.sentencepiece_model_file: '$DATA_DIR/tokenizers/sentencepiece/librispeech_bpe_1024.model'
 evalers['decoder_train'].metric_calculator.vocab.use_fast_tokenizer: False
 evalers['decoder_train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['decoder_train'].summary_writer.max_queue: 1000
 evalers['decoder_train'].summary_writer.write_every_n_steps: 1
 input.batcher.fn: 'axlearn.common.input_tf_data.batch'
 input.batcher.global_batch_size: 2048
@@ -561,4 +569,5 @@ prune_empty_state_updates: True
 save_input_iterator: False
 start_trace_process_indices[0]: 0
 summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+summary_writer.max_queue: 1000
 summary_writer.write_every_n_steps: 200

axlearn/experiments/testdata/axlearn.experiments.audio.conformer.librispeech_trainer/conformer-test-ctc.txt

@@ -250,4 +250,5 @@ prune_empty_state_updates: True
 save_input_iterator: False
 start_trace_process_indices[0]: 0
 summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+summary_writer.max_queue: 1000
 summary_writer.write_every_n_steps: 200

axlearn/experiments/testdata/axlearn.experiments.logistic_regression.tutorial/LogisticRegression.txt

@@ -18,6 +18,7 @@ evalers['eval'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAccu
 evalers['eval'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['eval'].metric_calculator.model_method: 'forward'
 evalers['eval'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['eval'].summary_writer.max_queue: 1000
 evalers['eval'].summary_writer.write_every_n_steps: 1
 input.klass: 'axlearn.common.input_grain.Input'
 input.source.fn: 'axlearn.experiments.logistic_regression.tutorial.build_source'
@@ -51,4 +52,5 @@ prune_empty_state_updates: True
 save_input_iterator: False
 start_trace_process_indices[0]: 0
 summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+summary_writer.max_queue: 1000
 summary_writer.write_every_n_steps: 10

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/envy-Switch-Base-single-host.txt

@@ -41,6 +41,7 @@ evalers['train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAcc
 evalers['train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['train'].metric_calculator.model_method: 'forward'
 evalers['train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['train'].summary_writer.max_queue: 1000
 evalers['train'].summary_writer.write_every_n_steps: 1
 evalers['validation'].eval_dtype: 'jax.numpy.bfloat16'
 evalers['validation'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -71,6 +72,7 @@ evalers['validation'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['validation'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['validation'].metric_calculator.model_method: 'forward'
 evalers['validation'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['validation'].summary_writer.max_queue: 1000
 evalers['validation'].summary_writer.write_every_n_steps: 1
 input.batcher.feed_batch_size: 8
 input.batcher.fn: 'axlearn.common.input_tf_data.per_feed_batch'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/envy-Switch-Base.txt

@@ -40,6 +40,7 @@ evalers['train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAcc
 evalers['train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['train'].metric_calculator.model_method: 'forward'
 evalers['train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['train'].summary_writer.max_queue: 1000
 evalers['train'].summary_writer.write_every_n_steps: 1
 evalers['validation'].eval_dtype: 'jax.numpy.bfloat16'
 evalers['validation'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -69,6 +70,7 @@ evalers['validation'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['validation'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['validation'].metric_calculator.model_method: 'forward'
 evalers['validation'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['validation'].summary_writer.max_queue: 1000
 evalers['validation'].summary_writer.write_every_n_steps: 1
 input.batcher.fn: 'axlearn.common.input_tf_data.per_feed_batch'
 input.batcher.pad_example_fn: 'axlearn.common.input_tf_data.default_pad_example_fn'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/envy-Switch-Large.txt

@@ -40,6 +40,7 @@ evalers['train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAcc
 evalers['train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['train'].metric_calculator.model_method: 'forward'
 evalers['train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['train'].summary_writer.max_queue: 1000
 evalers['train'].summary_writer.write_every_n_steps: 1
 evalers['validation'].eval_dtype: 'jax.numpy.bfloat16'
 evalers['validation'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -69,6 +70,7 @@ evalers['validation'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['validation'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['validation'].metric_calculator.model_method: 'forward'
 evalers['validation'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['validation'].summary_writer.max_queue: 1000
 evalers['validation'].summary_writer.write_every_n_steps: 1
 input.batcher.fn: 'axlearn.common.input_tf_data.per_feed_batch'
 input.batcher.pad_example_fn: 'axlearn.common.input_tf_data.default_pad_example_fn'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/envy-Switch-XXL.txt

@@ -40,6 +40,7 @@ evalers['train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAcc
 evalers['train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['train'].metric_calculator.model_method: 'forward'
 evalers['train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['train'].summary_writer.max_queue: 1000
 evalers['train'].summary_writer.write_every_n_steps: 1
 evalers['validation'].eval_dtype: 'jax.numpy.bfloat16'
 evalers['validation'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -69,6 +70,7 @@ evalers['validation'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['validation'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['validation'].metric_calculator.model_method: 'forward'
 evalers['validation'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['validation'].summary_writer.max_queue: 1000
 evalers['validation'].summary_writer.write_every_n_steps: 1
 input.batcher.fn: 'axlearn.common.input_tf_data.per_feed_batch'
 input.batcher.pad_example_fn: 'axlearn.common.input_tf_data.default_pad_example_fn'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/envy-test.txt

@@ -40,6 +40,7 @@ evalers['train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAcc
 evalers['train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['train'].metric_calculator.model_method: 'forward'
 evalers['train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['train'].summary_writer.max_queue: 1000
 evalers['train'].summary_writer.write_every_n_steps: 1
 evalers['validation'].eval_dtype: 'jax.numpy.bfloat16'
 evalers['validation'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -69,6 +70,7 @@ evalers['validation'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['validation'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['validation'].metric_calculator.model_method: 'forward'
 evalers['validation'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['validation'].summary_writer.max_queue: 1000
 evalers['validation'].summary_writer.write_every_n_steps: 1
 input.batcher.fn: 'axlearn.common.input_tf_data.per_feed_batch'
 input.batcher.pad_example_fn: 'axlearn.common.input_tf_data.default_pad_example_fn'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-fp8-single-host.txt

@@ -40,6 +40,7 @@ evalers['train'].metric_calculator.klass: 'axlearn.common.evaler.ModelSummaryAcc
 evalers['train'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['train'].metric_calculator.model_method: 'forward'
 evalers['train'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['train'].summary_writer.max_queue: 1000
 evalers['train'].summary_writer.write_every_n_steps: 1
 evalers['validation'].eval_dtype: 'jax.numpy.bfloat16'
 evalers['validation'].eval_policy.fn: 'axlearn.common.evaler.every_n_steps_policy'
@@ -69,6 +70,7 @@ evalers['validation'].metric_calculator.klass: 'axlearn.common.evaler.ModelSumma
 evalers['validation'].metric_calculator.metric_accumulator.klass: 'axlearn.common.metrics.MetricAccumulator'
 evalers['validation'].metric_calculator.model_method: 'forward'
 evalers['validation'].summary_writer.klass: 'axlearn.common.summary_writer.SummaryWriter'
+evalers['validation'].summary_writer.max_queue: 1000
 evalers['validation'].summary_writer.write_every_n_steps: 1
 input.batcher.fn: 'axlearn.common.input_tf_data.per_feed_batch'
 input.batcher.pad_example_fn: 'axlearn.common.input_tf_data.default_pad_example_fn'