Skip to content

Bridge checkpoint metrics into structured logging - #4189

Open
ivy-zhou wants to merge 1 commit into
pr4188from
pr4189
Open

Bridge checkpoint metrics into structured logging#4189
ivy-zhou wants to merge 1 commit into
pr4188from
pr4189

Conversation

@ivy-zhou

@ivy-zhou ivy-zhou commented Aug 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary:
torch_checkpointing emits EventLogger metadata on ordinary Python logging
records. TorchTitan kept its structured handlers on a separate non-propagating
logger, so those backend records reached the console but not the structured log.

Install one private handler on the root logger during
init_structured_logger(). It forwards records that carry log_type_name to
TorchTitan's existing structured handlers and ignores plain text. This matches
the llama4x logging model without a package-specific public forwarding API.

Async writers start in a fresh process, so preserve the resolved parent source,
output directory, and rank in the backend config. The child runs any existing
initializer first, restores TorchTitan structured logging, and admits backend
INFO events only when the backend logger is unset and inherits a stricter level.
Explicit DEBUG, WARNING, and ERROR settings remain unchanged.

Preserve the backend's context and measured_from_start_time_ms fields in
JSONL output, and stamp checkpoint_logging_context with the save step before
the backend snapshots it for asynchronous work.

Test Plan:
PYTHONPATH=/tmp/torchtitan-additional-packages.iGrJ91/checkpointing /home/ivyzhou/local/torchtitan-venv/bin/python -m pytest -q tests/unit_tests/observability/ tests/unit_tests/test_torch_checkpointing.py tests/unit_tests/test_checkpoint.py

162 passed.

PYTHONPATH=/tmp/torchtitan-additional-packages.iGrJ91/checkpointing:$PYTHONPATH PATH=/home/ivyzhou/local/torchtitan-venv/bin:$PATH /home/ivyzhou/local/torchtitan-venv/bin/pre-commit run --files tests/unit_tests/observability/test_structured_logging.py tests/unit_tests/test_torch_checkpointing.py torchtitan/components/checkpointer/torch_checkpointing.py torchtitan/observability/structured_logger/__init__.py torchtitan/observability/structured_logger/structured_logging.py

All hooks passed.

Stack created with Sapling. Best reviewed with ReviewStack.

This was referenced Aug 18, 2026
Open
Merged
Merged
Open
Merged
Merged
Open
Open
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 18, 2026
@ivy-zhou ivy-zhou mentioned this pull request Aug 18, 2026
Open
@ivy-zhou
ivy-zhou marked this pull request as ready for review August 18, 2026 15:00
ivy-zhou added a commit that referenced this pull request Aug 19, 2026
@ivy-zhou
Fix TorchFTCheckpointManager's contract with BaseCheckpointManager (#…
296624f 
…4183)

Pyrefly reports three errors in this file, all of them real:

_save was declared "-> None" while the base declares "-> bool", and it
discarded super()._save()'s result. BaseCheckpointManager.save returns
whatever
_save returns, so save() handed back None for this manager. Nothing
consumes it
today -- torchft/trainer.py ignores the result -- but the contract was
broken
and the next caller to check it would have been surprised. A replica
that skips
the full save now reports False; the per-replica dataloader checkpoint
is a side
channel, not the checkpoint this value describes.

_wait_for_saving dereferenced save_future without narrowing it. The
base's
maybe_wait_for_saving guarantees it is set before dispatching here,
which the
comment already said, so this just asserts what the comment claims.

_ft_save assigned dcp_save's "Future | AsyncSaveResponse | None"
straight into
save_future, typed "Future | None". AsyncMode.ASYNC always yields a
plain
Future, so assert that, matching how the DCP manager narrows the same
call in
its own ASYNC branch.

Only the first of the three is new: it arrived with the disabled-guard
refactor
(#4173), which renamed save to _save and made the base's return type
load
bearing. The other two predate it.

Test Plan:
  python3 -m pyrefly check torchtitan/experiments/torchft/checkpoint.py
  -> 0 errors (was 3)

python3 -m pytest
torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q
  -> 2 passed

Adds a test covering both branches of the participating_rank guard, so
the
return value is pinned rather than left to the type checker.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4183).
* #4191
* #4190
* #4189
* #4188
* #4197
* #4187
* #4186
* #4185
* #4184
* __->__ #4183
@ivy-zhou ivy-zhou mentioned this pull request Aug 19, 2026
Merged
@ivy-zhou ivy-zhou mentioned this pull request Aug 20, 2026
Merged
ivy-zhou added a commit that referenced this pull request Aug 20, 2026
@ivy-zhou
Remove the legacy torchtitan.components.checkpoint import path (#4184)
2e92be1 
Summary:
`components/checkpoint.py` is the last of the three re-export shims left
behind
when the checkpointer was grouped into a package, after the lr_scheduler
shim
in #4172 and checkpoint_utils in the preceding change. With it gone
there are
no compatibility shims left under `torchtitan/components/`.

This one is a pure module rename at the callsite.
`checkpointer/__init__.py` already re-exports exactly the same eight
symbols
the shim forwarded -- `AsyncMode`, `CheckpointManager`, `ModelWrapper`,
and the
`MODEL` / `OPTIMIZER` / `LR_SCHEDULER` / `DATALOADER` / `TRAIN_STATE`
key
constants -- so every importer changes only the module it names, with
the
imported names and their grouping untouched. Verified by parsing each
importer
and checking every imported symbol against the package's `__all__`
before
touching anything; nothing referenced a symbol the package does not
expose, and
no callsite used the plain `import torchtitan.components.checkpoint`
form.

Thirty modules are updated, spanning the checkpoint-conversion scripts,
the
forge, torchft, graph_trainer and rl experiments, and the unit tests. A
thirty-first file, `experiments/rl/__init__.py`, carries the import
inside its
module docstring as a usage example; that is updated too, so the
documented
path matches the working one.

This is an import-path change only; no runtime behavior changes.

Test Plan:
Full `pytest tests/unit_tests` (excluding `test_rope.py`, which cannot
be
collected without the optional `fla` package): 622 passed, 18 failed.
The 18
are the same set that fails on unmodified main in this environment --
missing
optional dependencies (`transformers`, `fla`) and environment-specific
kernel/compile failures (helion rope, inductor lora).

Also verified:
- No `components.checkpoint` references remain anywhere in the repo,
across all
  file types, and no `Compatibility imports` shim remains under
  `torchtitan/components/`.
- `import torchtitan.components.checkpoint` now raises
`ModuleNotFoundError`.
- All eight symbols import cleanly from
`torchtitan.components.checkpointer`.
- All 31 changed files byte-compile, and the affected non-test modules
(`trainer`, `forge.engine`, `torchft.checkpoint`, `torchft.optimizer`,
both
checkpoint-conversion scripts) import cleanly. `experiments.rl` fails
only on
  the absent optional `vllm` package, unrelated to this change.
- `ufmt` and `flake8 --config=.flake8` clean on all 31 files.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4184).
* #4240
* #4230
* #4191
* #4190
* #4189
* #4188
* #4197
* #4187
* #4186
* #4185
* __->__ #4184
ivy-zhou added a commit that referenced this pull request Aug 20, 2026
@ivy-zhou
Remove the legacy torchtitan.components.checkpoint_utils import path (#…
190e874 
…4185)

Summary:
`components/checkpoint_utils.py` became a re-export shim when the
optimizer and
checkpointer components were grouped into packages. Unlike the
lr_scheduler
shim removed in #4172, this one forwarded to two different destinations
at
once, which is what made it worth deleting rather than keeping: reading
an
import of `checkpoint_utils` told you nothing about whether the symbol
was
optimizer plumbing or checkpointer plumbing.

Route each of the four importers to the module that actually defines the
symbol. `canonical_fqn` lives in `checkpointer/utils.py`;
`init_optim_state`,
`get_flat_optim_state_dict`, and `load_flat_optim_state_dict` live in
`optimizer/utils.py`. This also settles the naming objection fegin
raised on
#4140, that `canonical_fqn` does not belong under an optimizer-shaped
name --
its importer in the rl trainer now names the checkpointer package
directly.

The three state-dict helpers are imported from `optimizer.utils` rather
than
re-exported through `optimizer/__init__.py`. They are low-level DCP
plumbing
with two callers between them, not part of the package's public surface,
which
stays `OptimizersContainer`, `LRSchedulersContainer`,
`ParamGroupConfig`, and
`default_adamw`.

This is an import-path change only; no runtime behavior changes.

Test Plan:
`pytest tests/unit_tests/test_state_dict_keys.py
tests/unit_tests/test_checkpoint.py
tests/unit_tests/test_lr_scheduler.py
tests/unit_tests/test_optimizer_param_groups.py
tests/unit_tests/test_torch_checkpointing.py`: 77 passed, 4 subtests
passed.

`test_legacy_checkpoint_utils_imports`, which asserted the shim's
symbols were
identical to the submodule's, is dropped -- it cannot outlive the shim.

`test_legacy_checkpoint_utils_can_be_imported_first` is kept but
retargeted, as
`test_state_dict_helpers_can_be_imported_first`. It guards a real
property
rather than the shim: the `optimizer` and `checkpointer` package
`__init__`
files import from each other, so importing either leaf `utils` module
first in
a fresh interpreter must not close an import cycle. It now subtests both
leaf
modules instead of the single shim.

Also verified:
- No `checkpoint_utils` references remain anywhere in the repo, across
all file
  types, not only Python.
- `import torchtitan.components.checkpoint_utils` now raises
  `ModuleNotFoundError`.
- `ufmt` and `flake8 --config=.flake8` clean on the four changed files.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4185).
* #4240
* #4230
* #4191
* #4190
* #4189
* #4188
* #4197
* #4187
* #4186
* __->__ #4185
* #4184
ivy-zhou
ivy-zhou commented Aug 21, 2026
View reviewed changes
Comment thread torchtitan/components/checkpointer/torch_checkpointing.py Outdated
# may have set this logger to DEBUG, and lowering it back to INFO would
# quietly discard the verbosity they asked for.
backend_logger = logging.getLogger(_BACKEND_LOGGER_NAME)
backend_logger.setLevel(min(backend_logger.getEffectiveLevel(), logging.INFO))

@ivy-zhou ivy-zhou Aug 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be clear, this doesn't change the root torchtitan logger, only sets the effective level for the torch_checkpointing logger to be a default INFO level because otherwise we wouldn't emit logs to the torchtitan JSONL side.

@tianyu-l
tianyu-l requested a review from felipemello1 August 21, 2026 18:50
tianyu-l
tianyu-l requested changes Aug 21, 2026
View reviewed changes

@tianyu-l tianyu-l left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the complexity doesn't smell right, would like @felipemello1 to take a look

@felipemello1

felipemello1 commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

the complexity doesn't smell right, would like @felipemello1 to take a look

I agree. I am looking into it. The TLDR is this:

  1. Theres another logger (checkpointer)
  2. We want to forward these logs to our SL (structured logger) API calls
  3. But then we also want to add 2 extra fields to it that are native from checkpointer

I am checking whats a better way of doing this

acisseJZhong pushed a commit that referenced this pull request Aug 21, 2026
@ivy-zhou @jessicazhongeee
Fix TorchFTCheckpointManager's contract with BaseCheckpointManager (#…
b838036 
…4183)

Pyrefly reports three errors in this file, all of them real:

_save was declared "-> None" while the base declares "-> bool", and it
discarded super()._save()'s result. BaseCheckpointManager.save returns
whatever
_save returns, so save() handed back None for this manager. Nothing
consumes it
today -- torchft/trainer.py ignores the result -- but the contract was
broken
and the next caller to check it would have been surprised. A replica
that skips
the full save now reports False; the per-replica dataloader checkpoint
is a side
channel, not the checkpoint this value describes.

_wait_for_saving dereferenced save_future without narrowing it. The
base's
maybe_wait_for_saving guarantees it is set before dispatching here,
which the
comment already said, so this just asserts what the comment claims.

_ft_save assigned dcp_save's "Future | AsyncSaveResponse | None"
straight into
save_future, typed "Future | None". AsyncMode.ASYNC always yields a
plain
Future, so assert that, matching how the DCP manager narrows the same
call in
its own ASYNC branch.

Only the first of the three is new: it arrived with the disabled-guard
refactor
(#4173), which renamed save to _save and made the base's return type
load
bearing. The other two predate it.

Test Plan:
  python3 -m pyrefly check torchtitan/experiments/torchft/checkpoint.py
  -> 0 errors (was 3)

python3 -m pytest
torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q
  -> 2 passed

Adds a test covering both branches of the participating_rank guard, so
the
return value is pinned rather than left to the type checker.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4183).
* #4191
* #4190
* #4189
* #4188
* #4197
* #4187
* #4186
* #4185
* #4184
* __->__ #4183
@ivy-zhou
Bridge checkpoint metrics into structured logging
fc7ff24 
Summary:
`torch_checkpointing` emits `EventLogger` metadata on ordinary Python logging
records. TorchTitan kept its structured handlers on a separate non-propagating
logger, so those backend records reached the console but not the structured log.

Install one private handler on the root logger during
`init_structured_logger()`. It forwards records that carry `log_type_name` to
TorchTitan's existing structured handlers and ignores plain text. This matches
the llama4x logging model without a package-specific public forwarding API.

Async writers start in a fresh process, so preserve the resolved parent source,
output directory, and rank in the backend config. The child runs any existing
initializer first, restores TorchTitan structured logging, and admits backend
INFO events only when the backend logger is unset and inherits a stricter level.
Explicit DEBUG, WARNING, and ERROR settings remain unchanged.

Preserve the backend's `context` and `measured_from_start_time_ms` fields in
JSONL output, and stamp `checkpoint_logging_context` with the save step before
the backend snapshots it for asynchronous work.

Test Plan:
`PYTHONPATH=/tmp/torchtitan-additional-packages.iGrJ91/checkpointing /home/ivyzhou/local/torchtitan-venv/bin/python -m pytest -q tests/unit_tests/observability/ tests/unit_tests/test_torch_checkpointing.py tests/unit_tests/test_checkpoint.py`

162 passed.

`PYTHONPATH=/tmp/torchtitan-additional-packages.iGrJ91/checkpointing:$PYTHONPATH PATH=/home/ivyzhou/local/torchtitan-venv/bin:$PATH /home/ivyzhou/local/torchtitan-venv/bin/pre-commit run --files tests/unit_tests/observability/test_structured_logging.py tests/unit_tests/test_torch_checkpointing.py torchtitan/components/checkpointer/torch_checkpointing.py torchtitan/observability/structured_logger/__init__.py torchtitan/observability/structured_logger/structured_logging.py`

All hooks passed.
This was referenced Aug 22, 2026
Open
Open
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tianyu-l tianyu-l tianyu-l requested changes
@fegin fegin Awaiting requested review from fegin fegin is a code owner
@wwwjn wwwjn Awaiting requested review from wwwjn wwwjn is a code owner
@wconstab wconstab Awaiting requested review from wconstab wconstab is a code owner
@felipemello1 felipemello1 Awaiting requested review from felipemello1

Assignees

No one assigned

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ivy-zhou @felipemello1 @tianyu-l