Multi gpu PyTorch autolog - creates multiple runs

[shrinath-suresh]

What changes are proposed in this pull request?

Fix #5817 .

When training in multiple gpu environment, PyTorch autolog logs model in multiple runs. The root cause of this problem - when autolog is invoked in multi gpu environment, the autolog scripts are running per gpu which in turn creates multiple runs.

Following are the possible solutions that i could think of

1. Wrapping patched_fit method with rank_zero_only decorator (PR changes)

Pros: Model will be logged only once.
Cons: Still multiple empty runs will be created during multi gpu training

2. Wrapping the autolog method with rank_zero_only decorator - click here to view the changes

Pros: Only one run will be created in multi gpu training - solves the root cause problem
Cons: creates dependency with pytorch-lightning import in mlflow pytorch library - need to evaluate if it will cause any import error (as of today pytorch lightning is imported only inside the autolog method)

3. Updating the examples to invoke autolog only with rank zero gpu - click here to view the changes

Pros: No library update needed
Cons: mlflow.pytorch.autolog() can be invoked only after instantiating the trainer object (as the trainer object is need to check the global rank)

MLFlow team and reviewers, please let us know your opinion.

How is this patch tested?

Autologging tests and existing unit tests

Does this PR change the documentation?

• No. You can skip the rest of this section.
• Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Check the status of the ci/circleci: build_doc check. If it's successful, proceed to the
   next step, otherwise fix it.
2. Click Details on the right to open the job page of CircleCI.
3. Click the Artifacts tab.
4. Click docs/build/html/index.html.
5. Find the changed pages / sections and make sure they render correctly.

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[BenWilson2]

Since a viable solution exists without introducing a core dependency change (like in option 2) and can avoid creating a different sort of issue (like in option 1), providing an update to examples (with appropriate notes mentioning why that block is in there) and a documentation update about how to configure autologging functionality for pytorch is probably a good idea.

If you wrap the block for this change with a conditional based on torch.cuda.device_count() and only trigger the alternative logic iif count > 1 (and, of course, that the training is in GPU mode i.e., cuda.is_available()) that should probably be pretty clear for people who are using multi-gpu pytorch.

Are you willing to update the examples and the documentation on pytorch?
Also, thank you for the thorough investigation into alternatives! It's very appreciated.

[shrinath-suresh]

Since a viable solution exists without introducing a core dependency change (like in option 2) and can avoid creating a different sort of issue (like in option 1), providing an update to examples (with appropriate notes mentioning why that block is in there) and a documentation update about how to configure autologging functionality for pytorch is probably a good idea.

If you wrap the block for this change with a conditional based on torch.cuda.device_count() and only trigger the alternative logic iif count > 1 (and, of course, that the training is in GPU mode i.e., cuda.is_available()) that should probably be pretty clear for people who are using multi-gpu pytorch.

Are you willing to update the examples and the documentation on pytorch? Also, thank you for the thorough investigation into alternatives! It's very appreciated.

@BenWilson2 Thanks for your comments. I have updated the MNIST example with the conditions for cpu and multi gpu training. Please let me know your thoughts. I will do it for all other examples if the changes looks fine.

[BenWilson2]

Added some notes to your MNIST example. Please feel free to file a PR with the changes there and in the other locations (including the main docs)!
Thanks!

[BenWilson2]

nit

Suggested change

	# In case of multi gpu training, the training script in invoked multiple times,
	# In case of multi gpu training, the training script is invoked multiple times,

[shrinath-suresh]

done

[BenWilson2]

nit

Suggested change

	# Since, the parameters are already logged using global rank 0 gpu, it is safe to ignore
	# Since the parameters are already logged using global rank 0 gpu, it is safe to ignore

[shrinath-suresh]

done

[BenWilson2]

A few nits to correct, please be sure to merge master into your branch to get the unit tests passing (the fix for the protobuf failures in examples tests is in master now), and be sure to run black on the examples.
It's looking good and very clear to users about how to prevent this unexpected behavior!

[BenWilson2]

LGTM! Thank you for the contribution!

[hikushalhere]

Hello @shrinath-suresh, I am confused how this change works. The DDP process group is initialized when PL's Trainer.fit() method is called (see here). Calling trainer.global_rank or torch.distributed.get_rank() before calling trainer.fit() throws the exception pasted below. What am I missing?

RuntimeError: Default process group has not been initialized, please make sure to call init_process_group.