feat: add production-ready MNIST example for PyTorch

[Snehadas2005]

What this PR does / why we need it:
This PR refactors and updates the existing PyTorch Jupyter Notebooks to fully support the Kubeflow Trainer V2 SDK. These updates transition the notebooks from legacy patterns to production-ready workflows that are compatible with the latest SDK features and cross-platform environments.

Updated Workflows:

• Image Classification (mnist.ipynb): Refactored to demonstrate the official V2 DDP training workflow on the Fashion MNIST dataset.
• Question Answering (fine-tune-distilbert.ipynb): Updated to demonstrate fine-tuning with Hugging Face integration, including critical fixes for offset mapping, Fast Tokenizers, and Accelerate backend requirements.
• Speech Recognition (speech-recognition.ipynb): Added a new transformer-based workflow for audio classification. Implements custom audio preprocessing/sampling using torchaudio and soundfile with native DDP support and multi-environment scaling.

Key Improvements:

• SDK V2 Native: Migrated all training logic to use TrainerClient and CustomTrainer.
• Windows & Local Compatibility: Standardized the distributed backend to gloo to ensure notebooks run successfully on Windows machines and local environments.
• Robust Verification Logic: Implemented a KUBEFLOW_TRAINER_TEST environment flag system. For LLM and Audio tasks, this now uses max_steps=1 and a tiny data subset to allow for instant logic verification without requiring high-compute resources.
• Unified Execution Model: Each notebook now provides three clear, tested paths for users:
• Direct Python: Quick kernel-level experimentation.
• SDK Local: Isolated environment verification using LocalProcessBackendConfig.
• Cluster Scaling: Distributed execution on Kubernetes with num_nodes scaling.
• Environment Stability: Added explicit dependency checks for accelerate, torchaudio, and tokenizers to prevent common runtime errors in notebook environments.

Which issue(s) this PR fixes:
Fixes #3062
Fixes #2040
PR: #2830

Checklist:

• Docs included if any changes are user facing (Updated PyTorch README)

[google-oss-prow]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign tenzen-y for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[jaiakash]

Thanks for raising this, @Snehadas2005.

I see its still a draft PR but few minor suggestions which will help you.

• Please use jupiter notebook (.ipynb) file for examples instead of standard python files (.py). You can checkout these examples PR (feat: qwen 2.5 1.5b runtime, example and fix gpu e2e test  #2835, feat(runtimes): Support Distributed MLX on CUDA #2790) as reference.
• You commits are not signed, thats why DCO is failing. Check this for more info on how to sign your current commits and even any of future commits.

Happy contributing.

[Snehadas2005]

Thank you so much, @jaiakash, for the detailed feedback and references. I really appreciate it.

That makes sense. I will convert the example into a Jupyter notebook and align it with the existing example patterns you shared, focusing on clarity and readability for data scientists.

I also appreciate the note on DCO signing. I will fix the commit signatures and ensure all future commits are properly signed.

Thanks again for the guidance, happy to iterate further and adjust based on feedback from the team.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB