change owner_params to owned_params

torch-ext/optimizer/muon.py

@@ -133,7 +133,7 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for v in per_dst), "all params should be sharded to all devices"
 
         send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
@@ -142,7 +142,7 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -176,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
 
             block = recv_counts[src]
             inner_off = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -255,7 +255,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
@@ -263,8 +263,8 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(