Repository navigation
fix(logging): use RANK env var in log_rank_zero before dist init - #2964
kratos0718 wants to merge 1 commit into
Conversation
Before torch.distributed is initialized (e.g. in recipe_main before setup() calls init_process_group), dist.is_initialized() returns False and the previous fallback was rank=0 unconditionally. This caused config.log_config — and any other log_rank_zero call made before distributed setup — to print on every rank in a torchrun launch. torchrun always sets the RANK environment variable before spawning worker processes, so we can use it as a reliable fallback to determine the true rank before the process group exists. Fixes meta-pytorch#2700 Signed-off-by: Abhinav Tarigoppula <[email protected]>
|
Hi @kratos0718! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at [email protected]. Thanks! |
|
̑ |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
meta-cla
Bot
added
the
CLA Signed
|
Hi! Following up — the Meta CLA is signed and all checks are passing. Happy to address any feedback or requested changes. Thanks! |
n-dlms
left a comment
n-dlms
left a comment
There was a problem hiding this comment.
Verified the root cause and the fix against current main. The chain is: recipes call config.log_config in recipe_main (for example full_finetune_distributed.py L1161) before init_process_group, config.log_config routes through log_rank_zero (config/_utils.py L28), and the primitive on main falls back to rank=0 whenever dist.is_initialized() is False (_logging.py L102). So under torchrun every rank is rank zero at that point and all of them log the full config. The RANK env fallback is the right fix: torchrun exports RANK before the process group exists, the pre-init path now reflects the real rank, and the no-RANK single-device case still logs.
Two observations from reading the rest of the call graph. First, this fixes more than the config log: log_rank_zero is called directly by a number of distributed recipes (lora_dpo, qat, full_dpo, lora_finetune, knowledge_distillation, full_finetune, qat_lora), so any pre-init call site inherits correct behavior with no per-recipe edits. Second, get_world_size_and_rank (training/_distributed.py L313) keeps its (1, 0) pre-init fallback, which is fine since recipes only call it after init_process_group, but it does mean this becomes the first env-based rank reader in the codebase; there is no existing helper to reuse and no duplication introduced.
The tests cover the three meaningful states (RANK=0 logs, RANK=1 stays silent, RANK absent logs) with dist mocked as available but not initialized, which matches the real pre-init condition. One edge examined rather than skipped: RANK set to an empty string makes int("") raise ValueError. torchrun never produces that, and failing loudly on a malformed hand-set RANK is better than silently defaulting to zero, so I would not change it.
From my side this is ready for approval, and it is the minimal root fix for #2700.
Fixes #2700
Summary
log_rank_zerofell back torank=0for all processes whendist.is_initialized()wasFalse(i.e. beforesetup()callsinit_process_group). This meantconfig.log_config— called inrecipe_mainbefore the recipe is constructed — printed the fullconfig on every rank in a multi-GPU torchrun launch.
torchrunalways setsRANKin the environment before spawningworkers, so we use it as the fallback rank when the process group is
not yet initialized.
What changed
torchtune/utils/_logging.py:log_rank_zeronow readsos.environ["RANK"](default"0") as fallback instead of assuming 0tests/torchtune/utils/test_logging.py: 3 new test cases covering pre-init RANK=0, RANK=1, and no-RANK (single device)Impact
Fixes redundant config logging in all distributed recipes
(
full_finetune_distributed,lora_finetune_distributed,knowledge_distillation_distributed, etc.) with a single change.