yoklainterview sim

Training Infrastructure ML Engineer Interview Questions

450 verified Training Infrastructure ML Engineer interview questions — solve with answers, learn from explanations, test yourself in a real simulation.

Try the real simulation →

Sample questions

Ti Checkpointing Fault ToleranceDifficulty 2
A team saves checkpoints with torch.save(model.state_dict(), path) only, then resumes training by loading that file into a freshly constructed model plus a freshly constructed SGD(momentum=0.9) optimizer. What breaks first after resume?
  • aThe model's forward pass produces NaN outputs immediately on the first resumed batch.
  • bThe optimizer has no momentum buffers yet, so early resumed steps rebuild momentum from zero.
  • cNothing breaks; SGD momentum is stored inside the model's parameters, not the optimizer.
  • dThe loss function silently switches to a different reduction mode after resume.
Explanation:A fresh SGD optimizer's state_dict()['state'] is empty until .step() has been called at least once, so momentum has to rebuild from scratch — the pre-crash momentum buffers were never saved because only model.state_dict() was persisted.
Ti Checkpointing Fault ToleranceDifficulty 2
A job trains with a StepLR scheduler that has decayed the learning rate down from 0.1 to 0.025 by step 12,000. Only model and optimizer state were saved. After a crash and resume with a brand-new StepLR instance, what learning rate does training resume at?
  • a0.025, because the optimizer's param_groups remembers the last applied learning rate independently of the scheduler.
  • b0.0125, because the scheduler halves the rate once more on the first resumed step.
  • c0.1, the scheduler's base learning rate, since a new StepLR starts with last_epoch=0.
  • d0.0, because an unsaved scheduler, lacking any restored state, defaults to a zero learning rate until re-warmed.
Explanation:A newly constructed StepLR always begins at last_epoch=0, so it applies the base learning rate passed to the optimizer's constructor. Since the scheduler's own state (its step counter) wasn't checkpointed, the run silently restarts the decay schedule from the beginning at the base rate.
Ti Checkpointing Fault ToleranceDifficulty 1
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
print(len(opt.state_dict()['state']))

This line runs right after opt is constructed, before any .step() call. What does it print?
  • a0
  • bThe number of parameter tensors in the model.
  • c1, since state holds one aggregate entry right after construction.
  • dAn error, because state_dict() cannot be called before the first optimization step.
Explanation:An optimizer's state dict is populated lazily — entries (like momentum buffers) are only created the first time a parameter is updated by .step(). Right after construction it's an empty dict, so len(...) is 0.
Ti Checkpointing Fault ToleranceDifficulty 1
Which of the following is actually stored inside torch.amp.GradScaler's state_dict()?
  • aThe full list of loss values observed since training started.
  • bA copy of every gradient tensor computed in the most recent backward pass.
  • cThe optimizer's learning rate for each parameter group.
  • dThe current loss-scale factor plus the growth/backoff factors and a growth-interval step counter.
Explanation:GradScaler.state_dict() returns a small dict such as {'scale', 'growth_factor', 'backoff_factor', 'growth_interval', '_growth_tracker'} — it tracks the current scaling factor and how close it is to the next scale-up, not gradients or loss history.
Ti Checkpointing Fault ToleranceDifficulty 2
A model checkpoint saves only weights and optimizer state, not the RNG state. After a crash, training resumes from the same step with a fresh torch.manual_seed(123) call at startup, hoping to reproduce the exact same dropout masks and augmentations as before the crash. What actually happens?
  • aThe RNG sequence is reproduced exactly, because seeding at startup reproduces the sequence from any later point in a run.
  • bThe RNG sequence diverges immediately, since re-seeding restarts the generator from zero.
  • cPyTorch automatically stores and restores the RNG state inside model.state_dict(), so this concern doesn't apply.
  • dDropout masks are deterministic regardless of RNG state once a model is in .train() mode.
Explanation:Re-seeding a generator with the same base seed replays the sequence from position zero, not from the point the pre-crash generator had reached after thousands of earlier draws — so the resumed run's dropout masks and augmentations no longer line up with the ones that would have followed pre-crash.
Ti Checkpointing Fault ToleranceDifficulty 2
A checkpoint-writing routine first calls torch.save(state, path + ".tmp") and, only after that completes without error, calls os.replace(path + ".tmp", path). Why write to a temporary path and rename, instead of calling torch.save(state, path) directly?
  • aIf the process dies mid-write, the temp file is left incomplete — the file at path stays untouched.
  • bIt's required because torch.save cannot overwrite an existing file at the same path.
  • cIt makes the write faster, since renaming a file — unlike a full rewrite — is quicker than writing the same number of bytes twice.
  • dIt's needed only when saving on GPU tensors; CPU tensors can always be saved directly to the final path without this risk.
Explanation:A rename (os.replace) on the same filesystem is effectively instantaneous and either fully succeeds or doesn't happen — so path always points to either the old, complete checkpoint or the new, complete one, never to a half-written file. Writing straight to path risks leaving a truncated, unreadable file there if the process dies mid-save.

Test yourself against the 2400-question ML Engineer bank.

Start interview