Skip to content

Simplify some tests since sharding propagation is in place#4304

Merged
Priya2698 merged 78 commits intomainfrom
wjy/prop
Apr 25, 2025
Merged

Simplify some tests since sharding propagation is in place#4304
Priya2698 merged 78 commits intomainfrom
wjy/prop

Conversation

@wujingyue
Copy link
Collaborator

No description provided.

@wujingyue wujingyue changed the base branch from main to pm/preseg_sharding_prop April 23, 2025 23:19
@wujingyue
Copy link
Collaborator Author

!test

@wujingyue wujingyue requested a review from Priya2698 April 23, 2025 23:19
@github-actions
Copy link

github-actions bot commented Apr 23, 2025

Review updated until commit 21834ae

Description

  • Simplify tests by removing redundant operations

  • Remove unnecessary sharding for output tensors

  • Update test cases to reflect changes


Changes walkthrough 📝

Relevant files
Tests
test_matmul.py
Simplify multidevice tests                                                             

tests/python/multidevice/test_matmul.py

  • Removed sharding and parallelization for output tensors in
    multidevice_schedule
  • Updated loops to exclude output tensors from device mesh setting
  • +2/-8     

    PR Reviewer Guide 🔍

    Here are some key observations to aid the review process:

    🧪 PR contains tests
    ⚡ Recommended focus areas for review

    Output Sharding

    The PR removes the sharding logic for the output tensor (self.out). This could lead to performance regressions or incorrect results if the output tensor was previously sharded for performance reasons.

    for t in [self.inp, self.weight, self.bias]:
        self.sched._set_device_mesh(t, mesh)
    
    # Shard N for weight (N, K) and bias (N)
    for t in [self.weight, self.bias]:
    Code Duplication

    The loop for setting the device mesh is duplicated in both multidevice_schedule methods. This could be refactored to avoid code duplication.

    for t in [self.inp, self.weight, self.bias]:
        self.sched._set_device_mesh(t, mesh)
    Tensor Definition

    The tensor definitions in definition methods have different shapes and biases. Ensure that these changes do not affect the intended functionality and performance of the tests.

    self.inp = self.define_tensor([-1, -1, d * e])
    self.weight = self.define_tensor([e, d * e])
    self.out = self.ops.linear(self.inp, self.weight, None)
    self.add_output(self.out)

    Base automatically changed from pm/preseg_sharding_prop to main April 24, 2025 03:43
    @wujingyue
    Copy link
    Collaborator Author

    !test

    @Priya2698 Priya2698 merged commit 87be6c3 into main Apr 25, 2025
    53 checks passed
    @Priya2698 Priya2698 deleted the wjy/prop branch April 25, 2025 19:13
    Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

    Labels

    None yet

    Projects

    None yet

    Development

    Successfully merging this pull request may close these issues.

    2 participants