Skip to content

Add --debug.print_config_ranks to limit config printing to select ranks - #4023

Open
githubsgi wants to merge 1 commit into
pytorch:mainfrom
githubsgi:sg/print-config-ranks
Open

Add --debug.print_config_ranks to limit config printing to select ranks#4023
githubsgi wants to merge 1 commit into
pytorch:mainfrom
githubsgi:sg/print-config-ranks

Conversation

@githubsgi

@githubsgi githubsgi commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds --debug.print_config_ranks, which restricts the --debug.print_config dump to a chosen set of global ranks.

Default is an empty list, meaning print on every rank — identical to today's behavior, so this is a no-op for existing users unless they opt in.

Motivation

logger in torchtitan is not rank-gated. Single-host runs stay readable because torchrun --local-ranks-filter ${LOG_RANK} drops non-selected ranks' stdout at the launcher.

That filtering isn't available everywhere. On our cluster jobs are launched with mpiexec -n N -ppn 8 --label, which forwards every rank's stdout into one combined log. --debug.print_config then writes the fully expanded config tree once per rank.

Measured on deepseek_v3_671b (n_layers=61), where the dump expands every per-layer sub-config:

Config dump, per rank 12,763 lines / 464,618 bytes
2048-rank run, unfiltered 8.02 GB total log, ~85% of it config dump
Same config at 256 nodes, with --debug.print_config_ranks 0 106 MB total, exactly 1 Running with configs block

Roughly a 10x reduction in log volume, and the log becomes greppable again.

This is generally useful to anyone driving torchtitan from a launcher that doesn't do rank filtering (MPI, Slurm srun without output redirection per task, and so on), rather than being specific to our setup.

Changes

  • torchtitan/config/configs.py — new print_config_ranks: list[int] field on DebugConfig.
  • torchtitan/config/manager.py — new list_int_rule tyro primitive rule so list[int] fields accept comma-separated values, mirroring the existing list_str_rule.
  • torchtitan/trainer.py — rank gate in maybe_log(). Reads the rank via torch.distributed.get_rank() when the process group is initialized, falling back to 0 otherwise. config.maybe_log() runs after init_distributed(), so the group is up in the distributed path.
  • tests/unit_tests/test_config_manager.py — covers the default and comma-separated parsing.

Usage

# unchanged default: every rank prints
--debug.print_config

# only rank 0 prints
--debug.print_config --debug.print_config_ranks 0

# a few ranks, e.g. one per parallelism group
--debug.print_config --debug.print_config_ranks 0,7,15

Testing

Parsing verified directly:

default         : []
explicit 0      : [0]
explicit 0,7,15 : [0, 7, 15]
empty           : []

End-to-end on a 256-node / 2048-rank deepseek_v3_671b job : with --debug.print_config_ranks 0 the log contains exactly one Running with configs block instead of 2048.

tests/unit_tests/test_config_manager.py::TestConfigManager::test_parse_print_config_ranks added and passing.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 30, 2026
@pytorch-bot

pytorch-bot Bot commented Jul 30, 2026

Copy link
Copy Markdown

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

  • ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant